import patch_gradio_jsonschema  # MUST be first - patches boolean JSON Schema handling
import patch_gradio_jsonschema  # MUST be first - patches boolean JSON Schema handling
import patch_gradio_jsonschema  # MUST be first - patches boolean JSON Schema handling\n
# app.py (Updated with Triage Orchestration)

import os
import math
from flask import Flask, render_template, request, jsonify
from dotenv import load_dotenv
from graph import triage_app, planner_app, main_app # Import all three compiled apps
from utils import generate_mermaid_diagram

load_dotenv()
app = Flask(__name__)

# Create necessary directories on startup
os.makedirs("outputs", exist_ok=True)
os.makedirs("uploads", exist_ok=True)
os.makedirs("memory", exist_ok=True)
os.makedirs("logs", exist_ok=True)


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


@app.route('/estimate', methods=['POST'])
def estimate():
    data = request.json
    user_input = data.get('message', '').strip()
    if not user_input:
        return jsonify({"error": "Message cannot be empty."}), 400
    
    # --- NEW TRIAGE STEP ---
    # First, check if the input is a simple greeting
    triage_inputs = {"userInput": user_input}
    triage_result = triage_app.invoke(triage_inputs)
    
    # If the triage agent provided a direct response, it's a greeting.
    if triage_result.get("draftResponse"):
        # We add a special key to let the frontend know to just display the message
        return jsonify({"is_greeting": True, "response": triage_result["draftResponse"]})

    # --- If not a greeting, proceed to the planner ---
    planner_inputs = {"userInput": user_input}
    try:
        estimate_result = planner_app.invoke(planner_inputs)
        estimate_result['pmPlan']['is_greeting'] = False
        return jsonify(estimate_result.get('pmPlan', {}))
    except Exception as e:
        return jsonify({"error": "An unexpected error occurred during planning."}), 500


@app.route('/chat', methods=['POST'])
def chat():
    data = request.json
    user_input = data.get('message', '').strip()
    
    try:
        user_budget = float(data.get('user_budget', 0.0))
        cost_per_loop = float(data.get('cost_per_loop', 0.05))
    except (ValueError, TypeError):
        return jsonify({"error": "Invalid budget or cost format."}), 400

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

    if cost_per_loop > 0:
        total_runs_affordable = max(1, math.floor(user_budget / cost_per_loop))
        max_loops_calibrated = total_runs_affordable - 1
    else:
        max_loops_calibrated = 0

    initial_state = {
        "userInput": user_input, "chatHistory": [], "coreObjectivePrompt": "",
        "retrievedMemory": "", "pmPlan": {}, "experimentCode": None,
        "experimentResults": None, "draftResponse": "", "qaFeedback": None,
        "approved": False, "execution_path": [], "rework_cycles": 0,
        "max_loops": max_loops_calibrated
    }
    
    try:
        final_state = main_app.invoke(initial_state)
    except Exception as e:
        return jsonify({"response": "An unexpected error occurred during execution. Please check the logs."}), 500
    
    run_path = final_state.get('execution_path', [])
    if run_path:
        mermaid_syntax = generate_mermaid_diagram(run_path)
        with open("outputs/last_run_flow.md", "w") as f:
            f.write("# Last Run Execution Flow\n\n")
            f.write("```mermaid\n")
            f.write(mermaid_syntax)
            f.write("```\n")
            
    response = final_state.get('draftResponse', "An error occurred, and no response was generated.")
    return jsonify({"response": response})

if __name__ == '__main__':
    app.run(debug=True, port=5001)