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('/upload', methods=['POST'])
def upload_file():
    """Handles file upload, saves it, and returns unique characters."""
    if 'file' not in request.files:
        return jsonify({"error": "No file part"}), 400
    file = request.files['file']
    if file.filename == '':
        return jsonify({"error": "No selected file"}), 400

    if file:
        session_id = get_session_id()
        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})

    return jsonify({"error": "File upload failed"}), 500

@app.route('/generate_map', methods=['POST'])
def generate_map():
    """Generates the initial wafer map data from the file and bin definitions."""
    session_id = get_session_id()
    if session_id not in session_data:
        return jsonify({"error": "Session not found. Please upload a file first."}), 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()}
    
    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 from file."}), 500

    # Store data in session
    session_data[session_id]['wafer_map'] = wafer_map.tolist()
    session_data[session_id]['char_to_bin_mapping'] = char_to_bin_mapping

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

@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')  # Expected format: [[row, col], [row, col], ...]
    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'])
    
    # Update the map based on the list of die coordinates
    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
            
    session_data[session_id]['wafer_map'] = wafer_map.tolist()

    return jsonify({"wafer_map": wafer_map.tolist()})

@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'])
    
    # np.rot90 rotates counter-clockwise, so k=-1 rotates clockwise
    rotated_map = np.rot90(wafer_map, k=-1)
    
    session_data[session_id]['wafer_map'] = rotated_map.tolist()

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

@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:
        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(debug=True)