"""
FastAPI app exposing the Git Conflict Resolution environment
as HTTP endpoints for HuggingFace Spaces + OpenEnv automated evaluation.
"""
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import Any
from environment import make_env, Action, GitConflictEnv

app = FastAPI(
    title="Git Conflict Resolution — OpenEnv",
    description="An OpenEnv environment where agents resolve Git merge conflicts.",
    version="1.0.0",
)

# In-memory session store (one env per task)
_envs: dict[str, GitConflictEnv] = {}


def get_env(task_id: str) -> GitConflictEnv:
    if task_id not in _envs:
        _envs[task_id] = make_env(task_id)
    return _envs[task_id]


# ---------------------------------------------------------------------------
# Required OpenEnv endpoints
# ---------------------------------------------------------------------------

@app.get("/")
def root():
    return {"status": "ok", "env": "git-conflict-resolver", "version": "1.0.0"}


@app.post("/reset")
def reset(task_id: str = "easy"):
    """Reset the environment and return the initial observation."""
    try:
        env = make_env(task_id)
        _envs[task_id] = env
        obs = env.reset()
        return obs.model_dump()
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))


@app.post("/step")
def step(task_id: str, action: dict[str, Any]):
    """Take a step: submit resolved files and get reward + observation."""
    env = get_env(task_id)
    try:
        act = Action(**action)
        obs, reward, done, info = env.step(act)
        return {
            "observation": obs.model_dump(),
            "reward": reward,
            "done": done,
            "info": info,
        }
    except RuntimeError as e:
        raise HTTPException(status_code=400, detail=str(e))
    except Exception as e:
        raise HTTPException(status_code=422, detail=str(e))


@app.get("/state")
def state(task_id: str = "easy"):
    """Return the current state of the environment."""
    env = get_env(task_id)
    return env.state()


@app.get("/tasks")
def list_tasks():
    """List all available tasks."""
    env = make_env("easy")
    return {"tasks": env.list_tasks()}


@app.get("/tasks/{task_id}/schema")
def task_schema(task_id: str):
    """Return the action schema for a specific task."""
    try:
        env = make_env(task_id)
        return env.get_task_schema(task_id)
    except ValueError as e:
        raise HTTPException(status_code=404, detail=str(e))


@app.post("/baseline")
def run_baseline_endpoint(task_id: str = "easy"):
    """
    Run baseline grader on the expected resolution to verify graders work.
    Used by automated evaluation to confirm scores are reproducible.
    """
    from tasks.scenarios import TASKS
    from graders.grader import grade

    if task_id not in TASKS:
        raise HTTPException(status_code=404, detail=f"Unknown task: {task_id}")

    expected = TASKS[task_id]["expected_resolution"]
    score, feedback = grade(task_id, expected)
    return {
        "task_id": task_id,
        "score": score,
        "feedback": feedback,
        "note": "Score on expected (ground-truth) resolution",
    }


@app.get("/grader")
def grader_score(task_id: str = "easy"):
    """Return the grader score after the last episode step."""
    env = get_env(task_id)
    s = env.state()
    return {
        "task_id": task_id,
        "last_score": s["last_score"],
        "last_feedback": s["last_feedback"],
        "attempts": s["attempts"],
        "done": s["done"],
    }

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

if __name__ == "__main__":
    main()