"""
SQL Query Optimizer — OpenEnv Environment Server
FastAPI server implementing the OpenEnv spec:
  POST /reset  → initial observation
  POST /step   → observation, reward, done, info
  GET  /state  → current state
"""

from __future__ import annotations

import re
import time
import uuid
from typing import Any

import sqlglot
import sqlglot.errors
from fastapi import FastAPI, HTTPException
from fastapi.responses import RedirectResponse
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field

from .tasks import TASKS, Task
from .grader import grade_submission

# ---------------------------------------------------------------------------
# Pydantic models (OpenEnv typed spec)
# ---------------------------------------------------------------------------

class ExecutionStats(BaseModel):
    estimated_cost: float = Field(description="Query planner estimated cost (lower = better)")
    scan_type: str = Field(description="Dominant scan type: seq_scan | index_scan | index_only_scan")
    rows_examined: int = Field(description="Estimated rows examined")
    has_anti_patterns: list[str] = Field(default_factory=list, description="Detected anti-patterns")


class Observation(BaseModel):
    task_id: str
    schema_ddl: str = Field(description="Table DDL describing the database schema")
    slow_query: str = Field(description="The SQL query to be optimized")
    execution_stats: ExecutionStats = Field(description="Stats for the slow (original) query")
    step_number: int = Field(default=0)
    last_reward: float | None = Field(default=None)
    last_feedback: str | None = Field(default=None)
    history: list[str] = Field(default_factory=list, description="Previous action summaries")
    done: bool = Field(default=False)


class Action(BaseModel):
    optimized_query: str = Field(description="Your rewritten, optimized SQL query")
    explanation: str = Field(default="", description="Brief explanation of changes made")


class Reward(BaseModel):
    total: float = Field(ge=0.0, le=1.0, description="Combined score 0.0–1.0")
    correctness: float = Field(ge=0.0, le=1.0, description="Does it return the same result set?")
    performance: float = Field(ge=0.0, le=1.0, description="Estimated cost reduction ratio")
    style: float = Field(ge=0.0, le=1.0, description="Absence of anti-patterns, readability")
    feedback: str = Field(description="Human-readable explanation of the score")


class StepResult(BaseModel):
    observation: Observation
    reward: Reward
    done: bool
    info: dict[str, Any] = Field(default_factory=dict)


class StateResponse(BaseModel):
    session_id: str
    task_id: str | None
    step_number: int
    total_reward: float
    best_reward: float
    is_done: bool


# ---------------------------------------------------------------------------
# In-memory session store
# ---------------------------------------------------------------------------

class Session:
    def __init__(self, task: Task) -> None:
        self.session_id = str(uuid.uuid4())
        self.task = task
        self.step_number = 0
        self.total_reward = 0.0
        self.best_reward = 0.0
        self.is_done = False
        self.history: list[str] = []
        self.created_at = time.time()

    def to_initial_observation(self) -> Observation:
        return Observation(
            task_id=self.task.task_id,
            schema_ddl=self.task.schema_ddl,
            slow_query=self.task.slow_query,
            execution_stats=ExecutionStats(**self.task.execution_stats),
            step_number=0,
            history=[],
            done=False,
        )

    def to_state(self) -> StateResponse:
        return StateResponse(
            session_id=self.session_id,
            task_id=self.task.task_id,
            step_number=self.step_number,
            total_reward=round(self.total_reward, 4),
            best_reward=round(self.best_reward, 4),
            is_done=self.is_done,
        )


_sessions: dict[str, Session] = {}
_current_session_id: str | None = None

# ---------------------------------------------------------------------------
# App
# ---------------------------------------------------------------------------

app = FastAPI(
    title="SQL Query Optimizer — OpenEnv",
    description="AI agent environment for learning to optimize SQL queries.",
    version="1.0.0",
)
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)


# ---------------------------------------------------------------------------
# Endpoints
# ---------------------------------------------------------------------------

@app.post("/reset", response_model=Observation, summary="Reset environment to a new task")
def reset(task_id: str | None = None) -> Observation:
    """
    Start a new episode. Optionally specify a task_id; otherwise the first
    task (easy) is selected. Returns the initial observation.
    """
    global _current_session_id

    if task_id is None:
        task_id = list(TASKS.keys())[0]

    if task_id not in TASKS:
        raise HTTPException(status_code=400, detail=f"Unknown task_id '{task_id}'. "
                            f"Valid: {list(TASKS.keys())}")

    task = TASKS[task_id]
    session = Session(task)
    _sessions[session.session_id] = session
    _current_session_id = session.session_id

    return session.to_initial_observation()


@app.post("/step", response_model=StepResult, summary="Submit an optimized query")
def step(action: Action) -> StepResult:
    """
    Submit the agent's optimized SQL query. Returns updated observation,
    reward breakdown, done flag, and debug info.
    """
    global _current_session_id

    if _current_session_id is None or _current_session_id not in _sessions:
        raise HTTPException(status_code=400, detail="No active session. Call /reset first.")

    session = _sessions[_current_session_id]

    if session.is_done:
        raise HTTPException(status_code=400, detail="Episode already done. Call /reset to start a new one.")

    session.step_number += 1
    MAX_STEPS = 5

    # Grade the submission
    reward_obj, grader_info = grade_submission(
        task=session.task,
        optimized_query=action.optimized_query,
        explanation=action.explanation,
    )

    session.total_reward += reward_obj.total
    session.best_reward = max(session.best_reward, reward_obj.total)

    # Record history
    summary = (
        f"Step {session.step_number}: score={reward_obj.total:.2f} "
        f"[correct={reward_obj.correctness:.2f}, perf={reward_obj.performance:.2f}, "
        f"style={reward_obj.style:.2f}]"
    )
    session.history.append(summary)

    # Episode ends on perfect score or max steps
    done = reward_obj.total >= 0.95 or session.step_number >= MAX_STEPS
    session.is_done = done

    obs = Observation(
        task_id=session.task.task_id,
        schema_ddl=session.task.schema_ddl,
        slow_query=session.task.slow_query,
        execution_stats=ExecutionStats(**session.task.execution_stats),
        step_number=session.step_number,
        last_reward=reward_obj.total,
        last_feedback=reward_obj.feedback,
        history=session.history.copy(),
        done=done,
    )

    return StepResult(
        observation=obs,
        reward=reward_obj,
        done=done,
        info={**grader_info, "session_id": session.session_id},
    )


@app.get("/state", response_model=StateResponse, summary="Current session state")
def state() -> StateResponse:
    """Return the current session state without advancing the episode."""
    if _current_session_id is None or _current_session_id not in _sessions:
        raise HTTPException(status_code=400, detail="No active session. Call /reset first.")
    return _sessions[_current_session_id].to_state()


@app.get("/tasks", summary="List all available tasks")
def list_tasks() -> dict:
    return {
        "tasks": [
            {
                "task_id": t.task_id,
                "difficulty": t.difficulty,
                "description": t.description,
            }
            for t in TASKS.values()
        ]
    }


@app.get("/health")
def health() -> dict:
    return {"status": "ok", "version": "1.0.0"}


@app.get("/", summary="Redirect to API docs", include_in_schema=False)
def root():
    return RedirectResponse(url="/docs")


def main():
    import uvicorn
    uvicorn.run("server.app:app", host="0.0.0.0", port=7860)

if __name__ == "__main__":
    main()