from dotenv import load_dotenv
load_dotenv()

from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse, HTMLResponse, RedirectResponse
from fastapi.middleware.cors import CORSMiddleware
from starlette.middleware.base import BaseHTTPMiddleware
from pydantic import BaseModel
from typing import Optional, List
import uvicorn
import os
from pathlib import Path
from environment.env import VoiceAuthenticityEnv


# ── Iframe-safe middleware for HF Spaces ────────────────────────────────
class HFSpacesMiddleware(BaseHTTPMiddleware):
    """Remove restrictive frame headers so HF Spaces iframe works."""
    async def dispatch(self, request: Request, call_next):
        response = await call_next(request)
        # Allow embedding in HF Spaces iframe
        response.headers["X-Frame-Options"] = "ALLOWALL"
        response.headers["Access-Control-Allow-Origin"] = "*"
        # Remove restrictive CSP if present
        try:
            del response.headers["content-security-policy"]
        except KeyError:
            pass
        return response


app = FastAPI(
    title="Voice Authenticity OpenEnv",
    description="Multi-step agentic environment for detecting synthetic speech",
    version="2.0.0"
)

# Add CORS middleware (needed for HF Spaces proxy)
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Add iframe-safe headers
app.add_middleware(HFSpacesMiddleware)

TASKS = [
    "clean_detection",
    "compressed_detection",
    "adversarial_detection",
    "streaming_detection",
    "phonecall_detection",
    "realtime_detection",
]

envs = {task: VoiceAuthenticityEnv(task) for task in TASKS}
current_task = "clean_detection"


class ActionRequest(BaseModel):
    action_type: str = "final_classify"
    label: int = 0
    confidence: float = 0.5
    reasoning: str = ""
    focus: List[str] = []
    task_name: Optional[str] = None


# Use __file__ to resolve Dashboard.html relative to THIS script,
# not os.getcwd() which can differ inside Docker containers.
_APP_DIR = Path(__file__).resolve().parent
DASHBOARD_PATH = _APP_DIR / "Dashboard.html"


def _load_dashboard() -> str:
    """Load dashboard HTML fresh every time to avoid stale caches in dev."""
    if not DASHBOARD_PATH.exists():
        raise FileNotFoundError(
            f"Dashboard.html not found at {DASHBOARD_PATH}"
        )
    return DASHBOARD_PATH.read_text(encoding="utf-8")


@app.get("/", response_class=HTMLResponse)
def root_interface():
    try:
        html = _load_dashboard()
        return HTMLResponse(
            content=html,
            headers={
                "X-Frame-Options": "ALLOWALL",
                "Content-Security-Policy": "frame-ancestors *",
            },
        )
    except FileNotFoundError as e:
        return HTMLResponse(
            content=f"<h2>Dashboard not found</h2><p>{e}</p>",
            status_code=404
        )


@app.get("/web", response_class=HTMLResponse)
def web_interface():
    try:
        html = _load_dashboard()
        return HTMLResponse(
            content=html,
            headers={
                "X-Frame-Options": "ALLOWALL",
                "Content-Security-Policy": "frame-ancestors *",
            },
        )
    except FileNotFoundError as e:
        return HTMLResponse(
            content=f"<h2>Dashboard not found</h2><p>{e}</p>",
            status_code=404
        )


@app.post("/reset")
def reset(request: dict = {}):
    global current_task
    task = request.get("task_name", current_task) if request else current_task
    if task not in envs:
        task = "clean_detection"
    current_task = task
    seed = request.get("seed") if request else None
    obs = envs[current_task].reset(seed=seed)
    return JSONResponse({
        "observation": obs.dict(),
        "done": False,
        "reward": 0.05,
        "info": {}
    })


@app.post("/step")
def step(action: ActionRequest):
    global current_task
    task = action.task_name or current_task
    if task not in envs:
        task = current_task
    action_dict = {
        "action_type": action.action_type,
        "label": action.label,
        "confidence": action.confidence,
        "reasoning": action.reasoning,
        "focus": action.focus,
    }
    obs, reward, done, info = envs[task].step(action_dict)
    return JSONResponse({
        "observation": obs.dict(),
        "reward": reward,
        "done": done,
        "info": info
    })


@app.get("/state")
def state():
    return JSONResponse(envs[current_task].state())


@app.get("/health")
def health():
    return {"status": "healthy", "service": "voice-authenticity-openenv"}


@app.get("/info")
def info():
    return {
        "name": "voice-authenticity-openenv",
        "version": "2.0.0",
        "status": "running",
        "tasks": TASKS,
        "web": "/web",
        "docs": "/docs"
    }


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


if __name__ == "__main__":
    main()