"""
FastAPI server for the SQL Debug Environment.
Exposes OpenEnv-compliant HTTP endpoints: /reset, /step, /state, /tasks, /grade
Includes optional web interface and WebSocket support for scaling.
"""

import os
import sys
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import uvicorn
import uuid
from typing import Optional
from fastapi import FastAPI, HTTPException, WebSocket, WebSocketDisconnect
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import HTMLResponse
from pydantic import BaseModel
from models import Action
from server.sql_environment import SQLDebugEnvironment
from tasks import TASKS, TASK_GROUPS
from graders import grade_submission


#Session Store
_environments: dict[str, SQLDebugEnvironment] = {}

ENABLE_WEB_INTERFACE = os.environ.get("ENABLE_WEB_INTERFACE", "false").lower() == "true"


def get_or_create_env(session_id: str) -> SQLDebugEnvironment:
    if session_id not in _environments:
        _environments[session_id] = SQLDebugEnvironment()
    return _environments[session_id]


#Response and Request Models
class ResetRequest(BaseModel):
    session_id: Optional[str] = None
    task_group: Optional[str] = None
    task_id: Optional[str] = None
    seed: Optional[int] = None


class StepRequest(BaseModel):
    session_id: str
    query: str
    explanation: Optional[str] = ""


class GradeRequest(BaseModel):
    task_id: str
    query: str

#FastAPI Config
app = FastAPI(
    title="SQL Debug Environment",
    description=(
        "An OpenEnv-compliant environment where AI agents learn to debug "
        "and fix broken SQL queries. Implements step() / reset() / state() API."
    ),
    version="1.0.0",
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)


#HTTP Endpoints
@app.get("/")
async def root():
    return {
        "name": "sql-debug-env",
        "version": "1.0.0",
        "description": "OpenEnv SQL Query Debugging Environment",
        "endpoints": ["/reset", "/step", "/state", "/tasks", "/grade", "/health", "/ws/{session_id}"],
        "task_groups": list(TASK_GROUPS.keys()),
        "web_interface": ENABLE_WEB_INTERFACE,
    }


@app.get("/health")
async def health():
    return {"status": "ok", "environment": "sql-debug-env", "active_sessions": len(_environments)}


@app.post("/reset")
async def reset(req: ResetRequest) -> dict:
    """Initialize or restart an episode. Returns initial observation."""
    session_id = req.session_id or str(uuid.uuid4())
    env = get_or_create_env(session_id)
    obs = env.reset(
        task_group=req.task_group,
        task_id=req.task_id,
        seed=req.seed,
    )
    return {
        "session_id": session_id,
        "observation": obs.model_dump(),
        "state": env.state().model_dump(),
    }


@app.post("/step")
async def step(req: StepRequest) -> dict:
    """Submit a SQL fix attempt. Returns observation with reward and done flag."""
    if req.session_id not in _environments:
        raise HTTPException(
            status_code=404,
            detail=f"Session '{req.session_id}' not found. Call /reset first.",
        )
    env = _environments[req.session_id]
    action = Action(query=req.query, explanation=req.explanation or "")
    obs = env.step(action)
    return {
        "session_id": req.session_id,
        "observation": obs.model_dump(),
        "reward": obs.reward,
        "done": obs.done,
        "success": obs.success,
        "state": env.state().model_dump(),
    }


@app.get("/state")
async def state(session_id: str) -> dict:
    """Return current episode state without taking an action."""
    if session_id not in _environments:
        raise HTTPException(status_code=404, detail="Session not found.")
    env = _environments[session_id]
    return env.state().model_dump()


@app.get("/tasks")
async def list_tasks() -> dict:
    """List all available tasks with metadata."""
    task_list = [
        {
            "id": task_id,
            "difficulty": task["difficulty"],
            "description": task["task_description"],
            "expected_columns": task["expected_columns"],
        }
        for task_id, task in TASKS.items()
    ]
    return {
        "tasks": task_list,
        "task_groups": {
            group: {"task_ids": ids, "count": len(ids)}
            for group, ids in TASK_GROUPS.items()
        },
    }


@app.post("/grade")
async def grade(req: GradeRequest) -> dict:
    """Grade a query against a specific task (no session needed)."""
    if req.task_id not in TASKS:
        raise HTTPException(status_code=404, detail=f"Task '{req.task_id}' not found.")
    score, feedback, details = grade_submission(req.task_id, req.query)
    return {"task_id": req.task_id, "score": score, "feedback": feedback, "details": details}

#WebSocket endpoint
@app.websocket("/ws/{session_id}")
async def websocket_endpoint(websocket: WebSocket, session_id: str):
    """
    WebSocket interface for real-time environment interaction.

    Message format (client → server):
        {"action": "reset", "task_group": "easy_syntax_fix", "seed": 42}
        {"action": "step",  "query": "SELECT ...", "explanation": "..."}
        {"action": "state"}

    Message format (server → client):
        {"type": "reset"|"step"|"state"|"error", ...observation/state fields...}
    """
    await websocket.accept()
    env = get_or_create_env(session_id)

    try:
        while True:
            data = await websocket.receive_json()
            action_type = data.get("action")

            if action_type == "reset":
                obs = env.reset(
                    task_group=data.get("task_group"),
                    task_id=data.get("task_id"),
                    seed=data.get("seed"),
                )
                await websocket.send_json({
                    "type": "reset",
                    "session_id": session_id,
                    "observation": obs.model_dump(),
                    "state": env.state().model_dump(),
                })

            elif action_type == "step":
                query = data.get("query", "")
                action = Action(query=query, explanation=data.get("explanation", ""))
                obs = env.step(action)
                await websocket.send_json({
                    "type": "step",
                    "session_id": session_id,
                    "observation": obs.model_dump(),
                    "reward": obs.reward,
                    "done": obs.done,
                    "success": obs.success,
                    "state": env.state().model_dump(),
                })

            elif action_type == "state":
                await websocket.send_json({
                    "type": "state",
                    "session_id": session_id,
                    "state": env.state().model_dump(),
                })

            else:
                await websocket.send_json({"type": "error", "detail": f"Unknown action: {action_type}"})

    except WebSocketDisconnect:
        pass

#Optional web interface
if ENABLE_WEB_INTERFACE:
    @app.get("/web", response_class=HTMLResponse)
    async def web_interface():
        """Interactive browser UI for debugging the environment manually."""
        return HTMLResponse(content=_WEB_UI_HTML)


_WEB_UI_HTML = """
<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<title>SQL Debug Env — Interactive UI</title>
<style>
  body { font-family: monospace; margin: 0; display: flex; height: 100vh; background: #1e1e1e; color: #d4d4d4; }
  #left, #right { flex: 1; padding: 1rem; overflow-y: auto; }
  #left { border-right: 1px solid #444; }
  h2 { color: #9cdcfe; }
  textarea { width: 100%; background: #252526; color: #d4d4d4; border: 1px solid #555; padding: 8px; font-family: monospace; }
  button { background: #0e639c; color: white; border: none; padding: 8px 16px; cursor: pointer; margin: 4px 2px; }
  button:hover { background: #1177bb; }
  select { background: #252526; color: #d4d4d4; border: 1px solid #555; padding: 6px; }
  pre { background: #252526; padding: 12px; white-space: pre-wrap; word-break: break-all; font-size: 12px; }
  .reward { color: #4ec9b0; font-weight: bold; }
  .error { color: #f44747; }
  .success { color: #4ec9b0; }
</style>
</head>
<body>
<div id="left">
  <h2>SQL Debug Environment</h2>
  <label>Task Group:
    <select id="taskGroup">
      <option value="easy_syntax_fix">Easy — Syntax Fix</option>
      <option value="medium_logic_fix">Medium — Logic Fix</option>
      <option value="hard_optimization_fix">Hard — Optimization Fix</option>
    </select>
  </label>
  <br><br>
  <button onclick="doReset()">Reset Episode</button>
  <hr>
  <h3>Your SQL Fix</h3>
  <textarea id="queryInput" rows="6" placeholder="Enter your corrected SQL query here..."></textarea>
  <br>
  <button onclick="doStep()">Submit Fix</button>
  <hr>
  <div id="feedback"></div>
</div>
<div id="right">
  <h2>Observation</h2>
  <pre id="obsOutput">Click "Reset Episode" to start.</pre>
</div>
<script>
  let sessionId = null;
  const base = window.location.origin;

  async function doReset() {
    const group = document.getElementById('taskGroup').value;
    const r = await fetch(base + '/reset', {
      method: 'POST',
      headers: {'Content-Type': 'application/json'},
      body: JSON.stringify({task_group: group})
    });
    const data = await r.json();
    sessionId = data.session_id;
    renderObs(data.observation);
    document.getElementById('feedback').innerHTML = '<em>Episode started. Submit your fix below.</em>';
    document.getElementById('queryInput').value = data.observation.buggy_query || '';
  }

  async function doStep() {
    if (!sessionId) { alert('Reset first!'); return; }
    const query = document.getElementById('queryInput').value;
    const r = await fetch(base + '/step', {
      method: 'POST',
      headers: {'Content-Type': 'application/json'},
      body: JSON.stringify({session_id: sessionId, query})
    });
    const data = await r.json();
    renderObs(data.observation);
    const cls = data.success ? 'success' : (data.reward > 0 ? 'reward' : 'error');
    document.getElementById('feedback').innerHTML =
      `<p class="${cls}">Reward: ${data.reward.toFixed(3)} | Done: ${data.done} | Success: ${data.success}</p>
       <p>${data.observation.last_attempt_result}</p>`;
  }

  function renderObs(obs) {
    document.getElementById('obsOutput').textContent = JSON.stringify(obs, null, 2);
  }
</script>
</body>
</html>
"""


if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)