"""
server.py
=========
FastAPI server — exposes InformationIntegrityEnv as an HTTP API.

Compatible with:
  - HuggingFace Spaces (gradio or docker)
  - Local Docker deployment
  - Any OpenEnv-compatible evaluation harness

Endpoints
---------
  POST /reset        Start a new episode
  POST /step         Submit one action
  GET  /state        Full internal state (debug)
  POST /collect      Collect a new scenario via DataCollector
  GET  /health       Liveness probe
  GET  /scenario     Return current scenario metadata (no ground truth)

Session management
------------------
  Each POST /reset returns a session_id.
  Subsequent /step calls MUST include session_id in the request body.
  Sessions are stored in Redis (TTL = SESSION_TTL_SECONDS).
  Fallback: in-process dict (single-worker mode).

Environment variables
---------------------
  REDIS_URL            (default redis://localhost:6379)
  SESSION_TTL_SECONDS  (default 3600)
  HOST                 (default 0.0.0.0)
  PORT                 (default 8000)
  LOG_LEVEL            (default info)
"""

import json
import logging
import os
import pickle
import traceback
import uuid
from contextlib import asynccontextmanager
from typing import Any, Dict, List, Optional

import uvicorn
from fastapi import Body, FastAPI, HTTPException, Query, Request, status
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, RedirectResponse
from pydantic import BaseModel, Field

from data_collector import CollectionRequest, DataCollector
from environment import InformationIntegrityEnv
from explainability import ExplainabilityEngine, build_reasoning_from_episode
from models import Action, ActionType, DecisionType, Observation

logger = logging.getLogger("server")

REDIS_URL         = os.getenv("REDIS_URL",           "redis://localhost:6379")
SESSION_TTL       = int(os.getenv("SESSION_TTL_SECONDS", 3600))
HOST              = os.getenv("HOST",                 "0.0.0.0")
PORT              = int(os.getenv("PORT",             7860))
LOG_LEVEL         = os.getenv("LOG_LEVEL",            "info")


# ---------------------------------------------------------------------------
# Session store (Redis-backed with in-process fallback)
# ---------------------------------------------------------------------------

class SessionStore:
    def __init__(self) -> None:
        self._mem: Dict[str, bytes] = {}
        self._redis = None
        try:
            import redis
            r = redis.from_url(REDIS_URL)
            r.ping()
            self._redis = r
            logger.info("Session store: Redis at %s", REDIS_URL)
        except Exception as exc:
            logger.warning("Redis unavailable (%s) — using in-process session store", exc)

    def get(self, session_id: str) -> Optional[InformationIntegrityEnv]:
        if self._redis:
            raw = self._redis.get(f"session:{session_id}")
        else:
            raw = self._mem.get(session_id)
        if raw is None:
            return None
        return pickle.loads(raw)

    def set(self, session_id: str, env: InformationIntegrityEnv) -> None:
        raw = pickle.dumps(env)
        if self._redis:
            self._redis.setex(f"session:{session_id}", SESSION_TTL, raw)
        else:
            self._mem[session_id] = raw

    def delete(self, session_id: str) -> None:
        if self._redis:
            self._redis.delete(f"session:{session_id}")
        else:
            self._mem.pop(session_id, None)

    def touch(self, session_id: str) -> None:
        if self._redis:
            self._redis.expire(f"session:{session_id}", SESSION_TTL)


_sessions = SessionStore()
_collector = DataCollector()
_explainer = ExplainabilityEngine()


# ---------------------------------------------------------------------------
# Request / response schemas
# ---------------------------------------------------------------------------

class ResetRequest(BaseModel):
    # ---- Programmatic difficulty control (OpenEnv eval path) ----
    difficulty: Optional[str] = Field(
        default=None,
        description="Load a deterministic scenario: 'easy' | 'medium' | 'hard'. "
                    "Mutually exclusive with `scenario`."
    )
    scenario_id: Optional[Any] = Field(
        default=None,
        description="Scenario index (int) or id (str) within the difficulty dataset. "
                    "None = random selection."
    )
    # ---- Legacy / custom scenario (DataCollector output) ----
    scenario: Optional[Dict[str, Any]] = Field(
        default=None,
        description="Pre-collected scenario dict (from POST /collect). "
                    "If provided, `difficulty` and `scenario_id` are ignored."
    )
    max_steps: int = Field(default=20, ge=1, le=100)


class ResetResponse(BaseModel):
    session_id: str
    observation: Dict[str, Any]
    message: str = "Episode started"


class StepRequest(BaseModel):
    session_id: str
    action: Dict[str, Any] = Field(
        ...,
        description='e.g. {"type":"verify_claim","claim_id":"claim_001"}'
    )


class StepResponse(BaseModel):
    session_id: str
    observation: Dict[str, Any]
    reward: Dict[str, Any]
    done: bool
    info: Dict[str, Any]
    explanation: Optional[str] = None


class CollectRequest(BaseModel):
    content: str
    claims: List[str] = Field(..., min_length=1)
    difficulty: str = "medium"
    topic: str = ""
    tags: List[str] = Field(default_factory=list)
    force_refresh: bool = False


class CollectResponse(BaseModel):
    scenario: Dict[str, Any]
    cached: bool = False


# ---------------------------------------------------------------------------
# App setup
# ---------------------------------------------------------------------------

@asynccontextmanager
async def lifespan(app: FastAPI):
    logging.basicConfig(
        level=getattr(logging, LOG_LEVEL.upper(), logging.INFO),
        format="%(asctime)s [%(levelname)s] %(name)s — %(message)s",
    )
    logger.info("InformationIntegrity env server starting on %s:%s", HOST, PORT)
    yield
    logger.info("Server shutting down")


app = FastAPI(
    title="InformationIntegrityEnv",
    description=(
        "OpenEnv-compatible deterministic benchmark environment for reasoning-based "
        "content moderation.\n\n"
        "Evaluation path: POST /reset with {\"difficulty\": \"easy|medium|hard\"} — "
        "no DataCollector or Redis required.\n\n"
        "Custom path: POST /collect then POST /reset with the returned scenario dict."
    ),
    version="1.1.0",
    lifespan=lifespan,
)

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


# ---------------------------------------------------------------------------
# Exception handler
# ---------------------------------------------------------------------------

@app.exception_handler(Exception)
async def global_exception_handler(request: Request, exc: Exception):
    logger.error("Unhandled exception: %s\n%s", exc, traceback.format_exc())
    return JSONResponse(
        status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
        content={"detail": str(exc)},
    )


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

@app.get("/")
async def root():
    return RedirectResponse(url="/ui/")


@app.get("/health")
async def health():
    return {"status": "healthy", "service": "InformationIntegrityEnv", "version": "1.1.0"}


@app.get("/metadata")
async def metadata():
    return {
        "name": "InformationIntegrityEnv",
        "description": (
            "A deterministic, benchmarkable content moderation environment where agents "
            "verify claims, reason over pre-built evidence sources, and make explainable "
            "moderation decisions (allow / flag / remove)."
        ),
        "version": "1.1.0",
    }


@app.get("/schema")
async def schema():
    return {
        "action": Action.model_json_schema(),
        "observation": Observation.model_json_schema(),
        "state": {
            "type": "object",
            "description": "Full internal environment state (debug only)",
            "properties": {
                "episode_id": {"type": "string"},
                "step_count": {"type": "integer"},
                "difficulty": {"type": "string"},
                "done": {"type": "boolean"},
                "ground_truth": {"type": "string"},
            },
        },
    }


@app.post("/mcp")
async def mcp(request: Request):
    try:
        body = await request.json()
    except Exception:
        body = {}
    req_id = body.get("id") if isinstance(body, dict) else None
    return {
        "jsonrpc": "2.0",
        "id": req_id,
        "result": {
            "tools": [],
            "description": "Use /reset, /step, /state directly.",
        },
    }


@app.get("/state")
async def get_state_root():
    """OpenAPI path discovery stub — returns empty state when no session specified."""
    return {"status": "no_active_session"}


@app.get("/tasks")

async def list_tasks():

    from tasks.graders import GraderFactory

    result = []

    for difficulty in ("easy", "medium", "hard"):

        try:

            grader = GraderFactory.get(difficulty)

            result.append({

                "id":          difficulty,

                "name":        difficulty,

                "difficulty":  difficulty,

                "has_grader":  True,

                "grader":      f"tasks.{difficulty}.grader.grade_episode",

                "grader_fn":   f"tasks.{difficulty}.grader.grade_episode",

                "score_range": [0.0, 1.0],

                "grader_weights": {

                    "decision":   grader.W_DECISION,

                    "evidence":   grader.W_EVIDENCE,

                    "reasoning":  grader.W_REASONING,

                    "efficiency": grader.W_EFFICIENCY,

                },

            })

        except Exception as exc:

            result.append({"id": difficulty, "error": str(exc), "has_grader": False})

    return result
 

@app.post("/collect", response_model=CollectResponse)
async def collect(req: CollectRequest):
    """
    Collect real-world evidence for a claim set and cache the scenario.
    Returns the scenario dict which can be passed directly to /reset.
    """
    request = CollectionRequest(
        content=req.content,
        claims=req.claims,
        difficulty=req.difficulty,
        topic=req.topic,
        tags=req.tags,
        force_refresh=req.force_refresh,
    )
    scenario = _collector.collect(request)
    return CollectResponse(scenario=scenario)


@app.post("/reset", response_model=ResetResponse)
async def reset(
    req: Optional[ResetRequest] = Body(default=None),
    difficulty: Optional[str] = Query(default=None),
    scenario_id: Optional[str] = Query(default=None),
):
    # Default to easy if no body provided OR if body is empty {}
    if req is None or (req.difficulty is None and req.scenario is None):
        req = ResetRequest(difficulty=difficulty or "easy", scenario_id=scenario_id)
    """
    Start a new episode.

    Query-parameter form (simplest — works with curl, browser, any HTTP client):
      POST /reset?difficulty=easy
      POST /reset?difficulty=medium&scenario_id=2
      POST /reset?difficulty=hard&scenario_id=hard_001

    JSON body form (same fields, body takes precedence):
      POST /reset  {"difficulty": "easy"}
      POST /reset  {"difficulty": "medium", "scenario_id": 2}
      POST /reset  {"scenario": <dict from /collect>}   ← legacy / live-data path
    """
    # Query params fill in body fields when the body fields are absent
    effective_difficulty  = req.difficulty  or difficulty
    effective_scenario_id = req.scenario_id if req.scenario_id is not None else scenario_id

    # Back-fill so the rest of the handler uses a single pair of variables
    req.difficulty  = effective_difficulty
    req.scenario_id = effective_scenario_id

    # Validate: at least one of difficulty or scenario must be provided
    if not req.difficulty and not req.scenario:
        raise HTTPException(
            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
            detail=(
                "Provide either 'difficulty' ('easy'|'medium'|'hard') to load a "
                "deterministic scenario, or 'scenario' (dict from POST /collect)."
            ),
        )

    if req.difficulty and req.difficulty not in ("easy", "medium", "hard"):
        raise HTTPException(
            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
            detail=f"Invalid difficulty '{req.difficulty}'. Must be easy, medium, or hard.",
        )

    session_id = str(uuid.uuid4())
    try:
        if req.scenario:
            # Legacy path — caller supplies a pre-collected scenario dict
            env = InformationIntegrityEnv(
                scenario=req.scenario,
                max_steps=req.max_steps,
            )
            obs = env.reset()
        else:
            # Eval path — load from static dataset (no DataCollector)
            env = InformationIntegrityEnv.from_dataset(
                difficulty=req.difficulty,
                scenario_id=req.scenario_id,
                max_steps=req.max_steps,
            )
            obs = env.reset()
    except (KeyError, IndexError, FileNotFoundError) as exc:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail=f"Scenario not found: {exc}",
        )
    except Exception as exc:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=f"Failed to create environment: {exc}",
        )

    _sessions.set(session_id, env)
    return ResetResponse(
        session_id=session_id,
        observation=obs.model_dump(),
    )


@app.post("/step", response_model=StepResponse)
async def step(req: StepRequest):
    """
    Submit one action and advance the episode.
    """
    env = _sessions.get(req.session_id)
    if env is None:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail=f"Session '{req.session_id}' not found or expired.",
        )

    try:
        action = Action(**req.action)
    except Exception as exc:
        raise HTTPException(
            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
            detail=f"Invalid action: {exc}",
        )

    try:
        obs, reward, done, info = env.step(action)
    except RuntimeError as exc:
        raise HTTPException(
            status_code=status.HTTP_409_CONFLICT,
            detail=str(exc),
        )

    # Persist updated env
    _sessions.set(req.session_id, env)
    _sessions.touch(req.session_id)

    # Build explanation and attach grader_score when episode ends
    explanation: Optional[str] = None
    if done:
        # Compute grader score and attach to info so the platform can verify it
        try:
            ep_result = env.episode_result()
            info["grader_score"] = ep_result.grader_score
            info["episode_result"] = ep_result.model_dump()
        except Exception as exc:
            logger.warning("grader_score computation failed: %s", exc)
            info.setdefault("grader_score", 0.005)

        try:
            decision = info.get("decision") or req.action.get("decision")
            reasoning = build_reasoning_from_episode(
                obs=obs,
                decision=decision or "none",
                episode_id=env._episode_id,
                difficulty=env.scenario.difficulty,
            )
            explanation = _explainer.generate_explanation(reasoning)
        except Exception as exc:
            logger.warning("Explanation generation failed: %s", exc)

    return StepResponse(
        session_id=req.session_id,
        observation=obs.model_dump(),
        reward=reward.model_dump(),
        done=done,
        info=info,
        explanation=explanation,
    )


@app.get("/state/{session_id}")
async def get_state(session_id: str):
    """Return full internal environment state (debug/audit endpoint)."""
    env = _sessions.get(session_id)
    if env is None:
        raise HTTPException(status_code=404, detail="Session not found")
    return env.state()


@app.get("/scenario/{session_id}")
async def get_scenario(session_id: str):
    """Return scenario metadata without ground truth (safe for agent use)."""
    env = _sessions.get(session_id)
    if env is None:
        raise HTTPException(status_code=404, detail="Session not found")
    sc = env.scenario
    return {
        "id":         sc.id,
        "difficulty": sc.difficulty,
        "topic":      sc.topic,
        "tags":       sc.tags,
        "content":    sc.content,
        "n_claims":   len(sc.claims),
        "n_sources":  len(sc.sources),
        "claim_ids":  [c.id for c in sc.claims],
        "source_ids": [s.id for s in sc.sources],
    }


@app.delete("/session/{session_id}")
async def delete_session(session_id: str):
    """Explicitly end a session and free memory."""
    _sessions.delete(session_id)
    return {"status": "deleted", "session_id": session_id}


class GradeRequest(BaseModel):
    """Request body for the /grade endpoint."""
    task_id: str = Field(..., description="Task difficulty: easy, medium, or hard")
    episode: Optional[dict] = Field(default=None, description="Episode data to grade")


@app.post("/grade")
async def grade_episode(req: GradeRequest):
    """
    Grade a completed episode for a specific task.
    Accepts an optional episode dict; if omitted, probes the grader with empty input.
    Returns {"score", "passed", "score_range", "grader"}.
    """
    import importlib

    task_id = req.task_id.lower().strip()
    if task_id not in ("easy", "medium", "hard"):
        raise HTTPException(status_code=400, detail=f"Unknown task_id '{task_id}'")

    try:
        mod   = importlib.import_module(f"tasks.{task_id}.grader")
        score = mod.grade_episode(req.episode or {})   # grade_episode(episode) -> float
        score = float(score)
        return {
            "task_id":     task_id,
            "score":       round(max(0.005, min(0.995, score)), 4),
            "passed":      True,
            "score_range": [0.0, 1.0],
            "grader":      f"tasks.{task_id}.grader.grade_episode",
        }
    except Exception as e:
        logger.error("Grading failed: %s", traceback.format_exc())
        raise HTTPException(status_code=500, detail=f"Grading failed: {str(e)}")


# ---------------------------------------------------------------------------
# Entrypoint
# ---------------------------------------------------------------------------