from flask import Flask, request, jsonify, render_template, send_from_directory, make_response
import os
import numpy as np
import wafer_processor as wp
from datetime import datetime

app = Flask(__name__)

# Configuration
UPLOAD_FOLDER = os.path.join('static', 'uploads')
OUTPUT_FOLDER = os.path.join('static', 'outputs')
app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
app.config['OUTPUT_FOLDER'] = OUTPUT_FOLDER

# Ensure directories exist
os.makedirs(UPLOAD_FOLDER, exist_ok=True)
os.makedirs(OUTPUT_FOLDER, exist_ok=True)

# In-memory storage for session data
session_data = {}

def get_session_id():
    """Gets or creates a unique session ID for the user."""
    if 'session_id' not in request.cookies:
        session_id = str(datetime.now().timestamp())
    else:
        session_id = request.cookies.get('session_id')
    return session_id

@app.route('/')
def index():
    """Serves the main HTML page and sets a session cookie."""
    session_id = get_session_id()
    resp = make_response(render_template('index.html'))
    resp.set_cookie('session_id', session_id)
    return resp

@app.route('/compare')
def compare_index():
    """Serves the comparison HTML page."""
    session_id = get_session_id()
    resp = make_response(render_template('compare.html'))
    resp.set_cookie('session_id', session_id)
    # Clean up previous comparison data if any
    if session_id in session_data and 'maps' in session_data[session_id]:
        del session_data[session_id]['maps']
    return resp

@app.route('/upload', methods=['POST'])
def upload_file():
    """
    Handles both single file upload for the editor and multiple file uploads for comparison.
    It distinguishes between them by checking the 'name' attribute of the file input.
    """
    session_id = get_session_id()

    # Case 1: Multiple files from comparison page (input name="files")
    if 'files' in request.files:
        files = request.files.getlist('files')
        if not files or files[0].filename == '':
            return jsonify({"error": "No files selected for comparison."}), 400

        session_data[session_id] = {'maps': {}}
        all_unique_chars = set()

        for i, file in enumerate(files):
            fid = f"map_{i}"
            filename = f"{session_id}_{fid}_{file.filename}"
            filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
            file.save(filepath)
            
            unique_chars = wp.get_unique_chars(filepath)
            if not unique_chars:
                continue

            all_unique_chars.update(unique_chars)
            session_data[session_id]['maps'][fid] = {
                'filepath': filepath,
                'filename': file.filename,
                'wafer_map': None,
                'reference': None
            }

        if not session_data[session_id]['maps']:
            return jsonify({"error": "None of the files could be processed."}), 400

        return jsonify({
            "unique_chars": sorted(list(all_unique_chars)),
            "maps": {fid: data['filename'] for fid, data in session_data[session_id]['maps'].items()}
        })

    # Case 2: Single file from main editor page (input name="file")
    elif 'file' in request.files:
        file = request.files['file']
        if file.filename == '':
            return jsonify({"error": "No selected file"}), 400

        filename = f"{session_id}_{file.filename}"
        filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
        file.save(filepath)

        unique_chars = wp.get_unique_chars(filepath)
        if not unique_chars:
            return jsonify({"error": "Could not process file. It might be empty or in the wrong format."}), 400

        session_data[session_id] = {'filepath': filepath}
        return jsonify({"unique_chars": unique_chars})

    # Case 3: No valid file part found
    return jsonify({"error": "No file part in the request."}), 400


@app.route('/generate_map', methods=['POST'])
def generate_map():
    """
    Generates wafer map data. For comparison mode, it generates for all uploaded maps.
    For editor mode, it also initializes the undo/redo history.
    """
    session_id = get_session_id()
    if session_id not in session_data:
        return jsonify({"error": "Session not found."}), 400

    data = request.json
    char_map = data.get('char_map')
    bin_type_to_value = {'good': 2, 'ng': 1, 'dummy': 0, 'ignore': -1}
    char_to_bin_mapping = {char: bin_type_to_value[bin_type] for char, bin_type in char_map.items()}

    # Comparison mode logic
    if 'maps' in session_data[session_id]:
        maps_data = {}
        for fid, map_info in session_data[session_id]['maps'].items():
            wafer_map = wp.read_wafer_map(map_info['filepath'], char_to_bin_mapping)
            if wafer_map.size == 0:
                continue
            session_data[session_id]['maps'][fid]['wafer_map'] = wafer_map.tolist()
            maps_data[fid] = {
                "wafer_map": wafer_map.tolist(),
                "rows": wafer_map.shape[0],
                "cols": wafer_map.shape[1]
            }
        session_data[session_id]['char_to_bin_mapping'] = char_to_bin_mapping
        return jsonify({"maps_data": maps_data})

    # Editor mode logic
    filepath = session_data[session_id]['filepath']
    wafer_map = wp.read_wafer_map(filepath, char_to_bin_mapping)
    if wafer_map.size == 0:
        return jsonify({"error": "Failed to read wafer map."}), 500

    wafer_map_list = wafer_map.tolist()
    session_data[session_id]['wafer_map'] = wafer_map_list
    session_data[session_id]['char_to_bin_mapping'] = char_to_bin_mapping
    
    # Initialize history for undo/redo
    session_data[session_id]['history'] = [wafer_map_list]
    session_data[session_id]['history_index'] = 0

    return jsonify({
        "wafer_map": wafer_map_list,
        "rows": wafer_map.shape[0],
        "cols": wafer_map.shape[1]
    })

def add_to_history(session_id, new_map_state):
    """Adds a new state to the history stack for undo/redo."""
    history = session_data[session_id].get('history', [])
    index = session_data[session_id].get('history_index', -1)
    
    # If we undo and then make a new change, truncate the future history
    if index < len(history) - 1:
        history = history[:index + 1]
        
    history.append(new_map_state)
    session_data[session_id]['history'] = history
    session_data[session_id]['history_index'] = len(history) - 1

@app.route('/set_reference', methods=['POST'])
def set_reference():
    """Sets the reference point for a specific map in comparison mode."""
    session_id = get_session_id()
    if session_id not in session_data or 'maps' not in session_data[session_id]:
        return jsonify({"error": "Not in comparison mode or session expired."}), 400

    data = request.json
    fid = data.get('fid')
    reference = data.get('reference') # Expected: [row, col]

    if fid not in session_data[session_id]['maps']:
        return jsonify({"error": "Invalid map ID."}), 400
    
    session_data[session_id]['maps'][fid]['reference'] = reference
    return jsonify({"success": True, "fid": fid, "reference": reference})

@app.route('/run_comparison', methods=['POST'])
def run_comparison():
    """Runs the comparison algorithm on all maps with set references."""
    session_id = get_session_id()
    if session_id not in session_data or 'maps' not in session_data[session_id]:
        return jsonify({"error": "Not in comparison mode or session expired."}), 400

    maps_to_compare = []
    references = []
    for fid, map_info in session_data[session_id]['maps'].items():
        if map_info.get('wafer_map') is not None and map_info.get('reference') is not None:
            maps_to_compare.append(np.array(map_info['wafer_map']))
            references.append(map_info['reference'])
        else:
            return jsonify({"error": f"Map '{map_info['filename']}' is missing data or a reference point."}), 400

    if len(maps_to_compare) < 2:
        return jsonify({"error": "At least two maps with reference points are required for comparison."}), 400

    comparison_result, legend, detailed_grid = wp.compare_wafer_maps(maps_to_compare, references)

    # Convert numpy types in detailed_grid to standard Python ints for JSON serialization
    serializable_grid = [[[int(b) for b in cell] for cell in row] for row in detailed_grid]

    # Store the detailed grid in the session for potential later use if needed
    session_data[session_id]['comparison_details'] = {
        'grid': serializable_grid,
        'map_names': [info['filename'] for info in session_data[session_id]['maps'].values() if info.get('wafer_map') is not None and info.get('reference') is not None]
    }

    return jsonify({
        "result_map": comparison_result.tolist(),
        "legend": legend,
        "detailed_grid": serializable_grid,
        "map_names": session_data[session_id]['comparison_details']['map_names'],
        "rows": comparison_result.shape[0],
        "cols": comparison_result.shape[1]
    })

@app.route('/compare/rotate', methods=['POST'])
def compare_rotate():
    """Rotates a single map within the comparison session."""
    session_id = get_session_id()
    if session_id not in session_data or 'maps' not in session_data[session_id]:
        return jsonify({"error": "Not in comparison mode or session expired."}), 400

    data = request.json
    fid = data.get('fid')

    if fid not in session_data[session_id]['maps']:
        return jsonify({"error": "Invalid map ID."}), 400
    
    map_info = session_data[session_id]['maps'][fid]
    wafer_map = np.array(map_info['wafer_map'])
    
    rotated_map = np.rot90(wafer_map, k=-1)
    map_info['wafer_map'] = rotated_map.tolist()
    
    # Clear reference if it was set, as rotation changes coordinates
    map_info['reference'] = None

    return jsonify({
        "fid": fid,
        "wafer_map": rotated_map.tolist(),
        "rows": rotated_map.shape[0],
        "cols": rotated_map.shape[1]
    })

@app.route('/compare/clear_reference', methods=['POST'])
def compare_clear_reference():
    """Clears the reference point for a specific map."""
    session_id = get_session_id()
    if session_id not in session_data or 'maps' not in session_data[session_id]:
        return jsonify({"error": "Not in comparison mode or session expired."}), 400

    data = request.json
    fid = data.get('fid')

    if fid not in session_data[session_id]['maps']:
        return jsonify({"error": "Invalid map ID."}), 400
        
    session_data[session_id]['maps'][fid]['reference'] = None
    
    return jsonify({"success": True, "fid": fid})

@app.route('/update_map', methods=['POST'])
def update_map():
    """Updates specified dies in the map and returns the full updated map."""
    session_id = get_session_id()
    if session_id not in session_data:
        return jsonify({"error": "Session expired or invalid."}), 400

    data = request.json
    dies_to_update = data.get('dies')
    new_bin_type = data.get('bin_type')
    
    if not dies_to_update:
        return jsonify({"error": "No dies specified for update."}), 400

    bin_type_to_value = {'good': 2, 'ng': 1, 'dummy': 0, 'ignore': -1}
    new_bin_value = bin_type_to_value[new_bin_type]

    wafer_map = np.array(session_data[session_id]['wafer_map'])
    
    for row, col in dies_to_update:
        if 0 <= row < wafer_map.shape[0] and 0 <= col < wafer_map.shape[1]:
            wafer_map[row, col] = new_bin_value
            
    wafer_map_list = wafer_map.tolist()
    session_data[session_id]['wafer_map'] = wafer_map_list
    add_to_history(session_id, wafer_map_list)

    return jsonify({"wafer_map": wafer_map_list})

@app.route('/rotate_map', methods=['POST'])
def rotate_map():
    """Rotates the current wafer map 90 degrees clockwise."""
    session_id = get_session_id()
    if session_id not in session_data:
        return jsonify({"error": "Session expired or invalid."}), 400

    wafer_map = np.array(session_data[session_id]['wafer_map'])
    rotated_map = np.rot90(wafer_map, k=-1)
    rotated_map_list = rotated_map.tolist()
    
    session_data[session_id]['wafer_map'] = rotated_map_list
    add_to_history(session_id, rotated_map_list)

    return jsonify({
        "wafer_map": rotated_map_list,
        "rows": rotated_map.shape[0],
        "cols": rotated_map.shape[1]
    })

@app.route('/inset_map', methods=['POST'])
def inset_map():
    """Applies an inset to the wafer map."""
    session_id = get_session_id()
    if session_id not in session_data:
        return jsonify({"error": "Session expired or invalid."}), 400

    data = request.json
    layers = int(data.get('layers', 1))
    from_bin_name = data.get('from_bin')
    to_bin_name = data.get('to_bin')

    bin_type_to_value = {'good': 2, 'ng': 1, 'dummy': 0, 'ignore': -1}
    from_bin_value = bin_type_to_value.get(from_bin_name)
    to_bin_value = bin_type_to_value.get(to_bin_name)

    if from_bin_value is None or to_bin_value is None:
        return jsonify({"error": "Invalid 'from' or 'to' bin type specified."}), 400

    wafer_map = np.array(session_data[session_id]['wafer_map'])
    inset_map_array = wp.inset_wafer_map(wafer_map, layers, from_bin_value, to_bin_value)
    inset_map_list = inset_map_array.tolist()
    
    session_data[session_id]['wafer_map'] = inset_map_list
    add_to_history(session_id, inset_map_list)

    return jsonify({
        "wafer_map": inset_map_list,
        "rows": inset_map_array.shape[0],
        "cols": inset_map_array.shape[1]
    })

@app.route('/undo', methods=['POST'])
def undo():
    session_id = get_session_id()
    if session_id not in session_data or not session_data[session_id].get('history'):
        return jsonify({"error": "No history to undo."}), 400
    
    history = session_data[session_id]['history']
    index = session_data[session_id]['history_index']

    if index > 0:
        index -= 1
        session_data[session_id]['history_index'] = index
        wafer_map = history[index]
        session_data[session_id]['wafer_map'] = wafer_map
        return jsonify({
            "wafer_map": wafer_map,
            "rows": len(wafer_map),
            "cols": len(wafer_map[0]) if wafer_map else 0,
            "can_undo": index > 0,
            "can_redo": True
        })
    return jsonify({"error": "Already at the beginning of history."}), 400

@app.route('/redo', methods=['POST'])
def redo():
    session_id = get_session_id()
    if session_id not in session_data or not session_data[session_id].get('history'):
        return jsonify({"error": "No history to redo."}), 400

    history = session_data[session_id]['history']
    index = session_data[session_id]['history_index']

    if index < len(history) - 1:
        index += 1
        session_data[session_id]['history_index'] = index
        wafer_map = history[index]
        session_data[session_id]['wafer_map'] = wafer_map
        return jsonify({
            "wafer_map": wafer_map,
            "rows": len(wafer_map),
            "cols": len(wafer_map[0]) if wafer_map else 0,
            "can_undo": True,
            "can_redo": index < len(history) - 1
        })
    return jsonify({"error": "Already at the end of history."}), 400

@app.route('/reset_map', methods=['POST'])
def reset_map():
    """Resets the map to its original state (the first state in history)."""
    session_id = get_session_id()
    if session_id not in session_data or not session_data[session_id].get('history'):
        return jsonify({"error": "No history available to reset from."}), 400
    
    history = session_data[session_id]['history']
    original_map = history[0]

    # Reset history to only the original state
    session_data[session_id]['history'] = [original_map]
    session_data[session_id]['history_index'] = 0
    session_data[session_id]['wafer_map'] = original_map

    return jsonify({
        "wafer_map": original_map,
        "rows": len(original_map),
        "cols": len(original_map[0]) if original_map else 0,
        "can_undo": False,
        "can_redo": False
    })

@app.route('/download', methods=['GET'])
def download_file():
    """Saves the current wafer map to a text file and serves it for download."""
    session_id = get_session_id()
    if session_id not in session_data or 'wafer_map' not in session_data[session_id]:
        return jsonify({"error": "No data to download."}), 400

    wafer_map = np.array(session_data[session_id]['wafer_map'])
    char_to_bin_mapping = session_data[session_id]['char_to_bin_mapping']
    bin_to_char_mapping = {v: k for k, v in char_to_bin_mapping.items()}

    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    output_filename = f"wafer_map_{timestamp}.txt"
    output_path = os.path.join(app.config['OUTPUT_FOLDER'], output_filename)

    wp.save_wafer_map(wafer_map, bin_to_char_mapping, output_path)

    return send_from_directory(app.config['OUTPUT_FOLDER'], output_filename, as_attachment=True)

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=12000, debug=True)