import os
import threading
import gc
import torch
from flask import Flask, request, jsonify, render_template, send_from_directory
from apartment_evaluator import ApartmentEvaluator

app = Flask(__name__, static_folder='static', template_folder='templates')

UPLOAD_FOLDER = os.path.abspath('uploads')
STATIC_STATIC_FOLDER = os.path.abspath('static')
CROPS_FOLDER = os.path.join(STATIC_STATIC_FOLDER, 'crops')

os.makedirs(UPLOAD_FOLDER, exist_ok=True)
os.makedirs(CROPS_FOLDER, exist_ok=True)

app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER

# Global variable to track the background job progress
PROGRESS = {
    "percent": 0,
    "status": "Idle",
    "running": False,
    "error": None,
    "result": None
}

evaluator_instance = None
evaluator_lock = threading.Lock()

def get_evaluator(max_size=480, fps=2.0):
    """Lazy initializer for the ApartmentEvaluator to avoid GPU memory overhead on startup."""
    global evaluator_instance
    with evaluator_lock:
        if evaluator_instance is None:
            PROGRESS["status"] = "Initializing SAM 2 and Grounding DINO models..."
            evaluator_instance = ApartmentEvaluator(max_size=max_size, fps=fps)
        else:
            evaluator_instance.max_size = max_size
            evaluator_instance.fps = fps
        return evaluator_instance


def run_evaluation_thread(before_path, after_path, prompt, output_dir, gemini_api_key=None, max_size=480, fps=2.0):
    global PROGRESS
    try:
        evaluator = get_evaluator(max_size=max_size, fps=fps)
        
        def update_progress(pct, desc):
            PROGRESS["percent"] = pct
            PROGRESS["status"] = desc

        result = evaluator.evaluate_apartment(
            before_video=before_path,
            after_video=after_path,
            prompt=prompt,
            output_dir=output_dir,
            gemini_api_key=gemini_api_key,
            progress_fn=update_progress
        )
        
        PROGRESS["result"] = result
        PROGRESS["percent"] = 100
        PROGRESS["status"] = "Complete!"
        PROGRESS["running"] = False
        
    except Exception as e:
        import traceback
        error_trace = traceback.format_exc()
        print("[Thread Error]", error_trace)
        PROGRESS["error"] = str(e)
        PROGRESS["running"] = False
        PROGRESS["status"] = f"Error: {str(e)}"
    finally:
        # Final cleanup
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        elif hasattr(torch, 'mps') and torch.backends.mps.is_available():
            torch.mps.empty_cache()


@app.route('/')
def index():
    return render_template('index.html')


@app.route('/upload', methods=['POST'])
def upload():
    global PROGRESS
    if PROGRESS["running"]:
        return jsonify({"error": "An evaluation is already running."}), 400

    if 'video_before' not in request.files or 'video_after' not in request.files:
        return jsonify({"error": "Missing video files."}), 400

    before_file = request.files['video_before']
    after_file = request.files['video_after']
    prompt = request.form.get('prompt', '').strip()

    if before_file.filename == '' or after_file.filename == '':
        return jsonify({"error": "No selected files."}), 400

    if not prompt:
        return jsonify({"error": "Prompt cannot be empty."}), 400

    # Save uploaded files
    before_path = os.path.join(app.config['UPLOAD_FOLDER'], 'before.mp4')
    after_path = os.path.join(app.config['UPLOAD_FOLDER'], 'after.mp4')
    
    before_file.save(before_path)
    after_file.save(after_path)

    # Reset progress status
    PROGRESS["percent"] = 0
    PROGRESS["status"] = "Videos uploaded successfully. Click Start to begin evaluation."
    PROGRESS["running"] = False
    PROGRESS["error"] = None
    PROGRESS["result"] = None

    return jsonify({"success": True, "prompt": prompt})


@app.route('/process', methods=['POST'])
def process():
    global PROGRESS
    if PROGRESS["running"]:
        return jsonify({"error": "Evaluation is already running."}), 400

    req_data = request.json or {}
    prompt = req_data.get('prompt', '').strip()
    gemini_api_key = req_data.get('gemini_api_key', '').strip() or None
    
    # Configure speed and frame rate
    speed_mode = req_data.get('speed_mode', 'fast') # 'super-fast' (360p), 'fast' (480p) or 'detailed' (720p)
    if speed_mode == 'super-fast':
        max_size = 360
    elif speed_mode == 'detailed':
        max_size = 720
    else:
        max_size = 480
    fps = float(req_data.get('fps', 2.0))

    if not prompt:
        return jsonify({"error": "Prompt is required."}), 400

    before_path = os.path.join(app.config['UPLOAD_FOLDER'], 'before.mp4')
    after_path = os.path.join(app.config['UPLOAD_FOLDER'], 'after.mp4')

    if not os.path.exists(before_path) or not os.path.exists(after_path):
        return jsonify({"error": "Uploaded video files not found. Upload them again."}), 400

    # Start tracking background thread
    PROGRESS["running"] = True
    PROGRESS["percent"] = 0
    PROGRESS["status"] = "Starting analysis job..."
    PROGRESS["error"] = None
    PROGRESS["result"] = None

    thread = threading.Thread(
        target=run_evaluation_thread,
        args=(before_path, after_path, prompt, STATIC_STATIC_FOLDER, gemini_api_key, max_size, fps)
    )
    thread.daemon = True
    thread.start()

    return jsonify({"success": True})


@app.route('/status', methods=['GET'])
def status():
    return jsonify(PROGRESS)


@app.route('/reset', methods=['POST'])
def reset():
    global PROGRESS
    req_data = request.json or {}
    force = req_data.get('force', False)
    
    if PROGRESS["running"] and not force:
        return jsonify({"error": "Cannot reset while evaluation is running. Use force=true to override."}), 400
    
    PROGRESS["percent"] = 0
    PROGRESS["status"] = "Idle"
    PROGRESS["running"] = False
    PROGRESS["error"] = None
    PROGRESS["result"] = None
    return jsonify({"success": True})


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