# server/app.py
# FastAPI server that wraps the StudyEnvironment and exposes it as an HTTP API
#
# Run with: uvicorn server.app:app --host 0.0.0.0 --port 8000
# or just: python -m uvicorn server.app:app --reload (for dev)

import sys
import os

# make sure we can import from src/ regardless of how this is run
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import Optional
import uvicorn

from src.environment import StudyEnvironment
from src.models import Action, StepResult, Observation

app = FastAPI(
    title="Smart Study Session Optimizer",
    description="An RL environment for optimizing study sessions. Part of OpenEnv hackathon.",
    version="1.0.0",
)

# cors - needed for HuggingFace spaces and browser clients
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

# global environment instance - one env per server process
# this is fine for single-agent use, would need sessions for multi-agent
env: Optional[StudyEnvironment] = None

print("[SERVER] Study Optimizer server starting up...")


class ResetRequest(BaseModel):
    task_id: Optional[str] = "maximize_productivity"


class StepRequest(BaseModel):
    action: int


@app.get("/health")
def health():
    """Health check. Evaluators ping this first."""
    return {"status": "ok"}


@app.post("/reset")
def reset(request: ResetRequest = ResetRequest()):
    """
    Reset the environment. Call this before starting a new episode.
    Returns the first observation.
    """
    global env

    task = request.task_id or "maximize_productivity"
    print(f"[SERVER] /reset called, task_id={task}")

    env = StudyEnvironment(task_id=task)
    obs = env.reset()

    # return as dict so it plays nice with the OpenEnv spec
    return {
        "observation": obs.model_dump(),
        "task_id": task,
        "info": {
            "score": 0.5,
            "task_id": task,
            "step_count": 0,
            "all_task_scores": {
                "survive_8_hours": 0.5,
                "maximize_productivity": 0.5,
                "smart_break_taker": 0.5,
            }
        }
    }


def _clamp_score(v: float) -> float:
    """Last-resort clamp: force any score into [0.1, 0.9]."""
    try:
        val = float(v)
        return max(0.1, min(0.9, val))
    except:
        return 0.5


@app.post("/step")
def step(request: StepRequest):
    """
    Execute one step in the environment.
    Returns observation, reward, done, info.
    """
    global env

    if env is None:
        # auto-init with defaults if reset wasn't called first
        print("[SERVER] WARNING: /step called before /reset, auto-initializing")
        env = StudyEnvironment()
        env.reset()

    if not 0 <= request.action <= 2:
        raise HTTPException(status_code=400, detail=f"Invalid action {request.action}. Must be 0, 1, or 2.")

    print(f"[SERVER] /step called, action={request.action}")
    result: StepResult = env.step(request.action)

    info = result.info

    # SAFETY NET: clamp all scores at the API boundary
    if "score" in info:
        info["score"] = _clamp_score(info["score"])
    if "all_task_scores" in info:
        for task_key in info["all_task_scores"]:
            info["all_task_scores"][task_key] = _clamp_score(info["all_task_scores"][task_key])

    return {
        "observation": result.observation.model_dump(),
        "reward": result.reward,
        "done": result.done,
        "info": info,
    }


@app.get("/state")
def state():
    """Return the full internal environment state. Useful for debugging."""
    global env

    if env is None:
        return {"status": "not_initialized", "message": "Call /reset first"}

    return env.get_state()


@app.get("/")
def root():
    """Basic info page."""
    return {
        "name": "Smart Study Session Optimizer",
        "version": "1.0.0",
        "description": "RL environment for optimizing study sessions",
        "endpoints": ["/health", "/reset", "/step", "/state"],
        "tasks": ["survive_8_hours", "maximize_productivity", "smart_break_taker"],
    }


def main():
    """Entry point for the server."""
    port = int(os.environ.get("PORT", 7860))
    print(f"[SERVER] Starting on port {port}")
    uvicorn.run(app, host="0.0.0.0", port=port)


if __name__ == "__main__":
    main()