from fastapi import HTTPException
from pydantic import BaseModel
from typing import Dict

from openenv.core import create_app

from models import Observation, Action
from env.environment import SQLEnv
from env.tasks import TASKS

app = create_app(
    env=SQLEnv,
    action_cls=Action,
    observation_cls=Observation,
    env_name="sql-query-optimizer",
)


@app.get("/tasks")
async def get_tasks():
    action_schema = Action.model_json_schema()
    task_list = [{"id": k, **v} for k, v in TASKS.items()]
    return {
        "tasks": task_list,
        "action_schema": action_schema,
    }


class BaselineResponse(BaseModel):
    scores: Dict[int, float]


@app.post("/baseline", response_model=BaselineResponse)
async def run_baseline():
    import baseline

    try:
        scores = baseline.run_all_tasks()
        return BaselineResponse(scores=scores)
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


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


if __name__ == "__main__":
    main()