import os
import uuid
import json
import time
import requests
from pathlib import Path
from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse, HTMLResponse
from fastapi.templating import Jinja2Templates

# =========================
# APP CORE
# =========================

app = FastAPI(title="PyRunner PAO v5 SaaS Core", version="5.0")

BASE_DIR = Path(__file__).resolve().parent

# SAFE TEMPLATE LOADING (prevents your Jinja crash)
TEMPLATE_DIR = BASE_DIR / "templates"
templates = Jinja2Templates(directory=str(TEMPLATE_DIR))

# =========================
# CONFIG
# =========================

OLLAMA_URL = os.getenv("OLLAMA_URL", "http://localhost:11434")

MODELS = {
    "planner": "qwen2.5:1.5b",
    "worker": "qwen2.5-coder:1.5b",
    "critic": "deepseek-coder:1.3b",
    "synth": "llama3.2:1b"
}

# =========================
# MEMORY SYSTEM (replace later with Supabase/Postgres)
# =========================

MEMORY = {}

def memory_store(task_id, data):
    MEMORY[task_id] = data

def memory_get(task_id):
    return MEMORY.get(task_id, {})

# =========================
# EVENT SYSTEM (lightweight observability)
# =========================

EVENTS = []

def emit(event_type, payload):
    EVENTS.append({
        "id": str(uuid.uuid4()),
        "type": event_type,
        "payload": payload,
        "ts": time.time()
    })

def drain_events():
    data = EVENTS[:]
    EVENTS.clear()
    return data

# =========================
# LLM CALL
# =========================

def call_llm(prompt, model):
    try:
        r = requests.post(
            f"{OLLAMA_URL}/api/generate",
            json={
                "model": model,
                "prompt": prompt,
                "stream": False
            },
            timeout=120
        )
        return r.json().get("response", "")
    except Exception as e:
        return f"LLM_ERROR: {str(e)}"

# =========================
# DAG PLANNER (FIXED JSON SAFETY)
# =========================

def plan_task(task: str):

    prompt = f"""
Return ONLY valid JSON.

Break task into steps:

Task: {task}

Format:
[
  "step 1",
  "step 2",
  "step 3"
]
"""

    raw = call_llm(prompt, MODELS["planner"])

    try:
        steps = json.loads(raw)

        # ensure list safety
        if isinstance(steps, list):
            return [{"step": s} for s in steps]

    except:
        pass

    # fallback (always safe)
    return [{"step": task}]

# =========================
# TOOL SYSTEM
# =========================

def tool_calculator(expr: str):
    try:
        return eval(expr, {"__builtins__": {}})
    except:
        return "calc_error"

TOOLS = {
    "calculator": tool_calculator
}

def run_tools(text: str):
    if "calc:" in text:
        expr = text.replace("calc:", "").strip()
        return str(TOOLS["calculator"](expr))
    return text

# =========================
# WORKER NODE
# =========================

def worker(step: str, context: str):

    prompt = f"""
Execute this instruction clearly:

Step: {step}

Context:
{context}
"""

    return call_llm(prompt, MODELS["worker"])

# =========================
# CRITIC (QUALITY CHECK)
# =========================

def critic(task, output):

    prompt = f"""
Rate and improve:

Task: {task}
Output: {output}

Return short improvement suggestion.
"""

    return call_llm(prompt, MODELS["critic"])

# =========================
# SYNTHESIZER
# =========================

def synth(outputs):

    combined = "\n".join(outputs)

    prompt = f"""
Create final clean response:

{combined}
"""

    return call_llm(prompt, MODELS["synth"])

# =========================
# EXECUTION ENGINE
# =========================

def run_dag(task: str):

    task_id = str(uuid.uuid4())

    emit("task_started", {"task_id": task_id, "task": task})

    dag = plan_task(task)
    results = []

    context = task

    for node in dag:

        step = node["step"]

        emit("step_started", {"step": step})

        raw = worker(step, context)
        processed = run_tools(raw)

        _crit = critic(task, processed)

        results.append(processed)
        context = processed

        emit("step_done", {
            "step": step,
            "output": processed,
            "critique": _crit
        })

    final = synth(results)

    memory_store(task_id, {
        "task": task,
        "result": final,
        "steps": results
    })

    emit("task_done", {"task_id": task_id})

    return {
        "task_id": task_id,
        "dag": dag,
        "events": drain_events(),
        "final": final
    }

# =========================
# API ROUTES
# =========================

@app.get("/", response_class=HTMLResponse)
async def home(request: Request):

    # SAFE CHECK: prevents template crash
    if not (TEMPLATE_DIR / "index.html").exists():
        return HTMLResponse("<h1>Missing index.html</h1>", status_code=500)

    return templates.TemplateResponse("index.html", {"request": request})


@app.get("/builder", response_class=HTMLResponse)
async def builder(request: Request):

    if not (TEMPLATE_DIR / "builder.html").exists():
        return HTMLResponse("<h1>Missing builder.html</h1>", status_code=500)

    return templates.TemplateResponse("builder.html", {"request": request})


@app.post("/api/run")
async def run(request: Request):

    body = await request.json()
    task = body.get("task", "")

    if not task:
        return JSONResponse({"error": "missing task"}, status_code=400)

    return run_dag(task)


@app.get("/api/memory/{task_id}")
def get_memory(task_id: str):
    return memory_get(task_id)


@app.get("/api/events")
def get_events():
    return {"events": EVENTS}