from fastapi import FastAPI, APIRouter, HTTPException
from fastapi.responses import HTMLResponse, FileResponse, RedirectResponse
from fastapi.staticfiles import StaticFiles
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import Optional
import os

app = FastAPI(title="ARIA — Autonomous Reasoning & Intelligence Architecture")

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


@app.get("/health")
async def health():
    return {"status": "healthy", "system": "ARIA"}


@app.get("/metadata")
async def metadata():
    return {
        "project": "ARIA",
        "full_name": "Autonomous Reasoning & Intelligence Architecture",
        "version": "1.0.0",
        "hackathon": "FAR AWAY 2026",
        "theme": "Agentic & Autonomous Systems",
        "description": "Universal autonomous multi-agent orchestration system",
    }


# Mount static files
if not os.path.exists("ui/static"):
    os.makedirs("ui/static", exist_ok=True)
app.mount("/static", StaticFiles(directory="ui/static"), name="static")


# ─────────────────────────────────────────────
# ARIA Orchestration Routes
# ─────────────────────────────────────────────

from pathlib import Path as _Path

ARIA_UI_FILE = _Path(__file__).resolve().parent / "ui" / "gradio_demo.py"
_aria_env = None


class ARIARunRequest(BaseModel):
    goal: str
    task_id: str = "medium"
    seed: int = 42


class ARIAStepRequest(BaseModel):
    action_type: str
    agent_id: str
    reason: Optional[str] = None


class ARIAPlanRequest(BaseModel):
    agent_id: str
    assigned_task: str
    priority: int = 3
    reason: Optional[str] = None


@app.post("/aria/run")
async def aria_run(req: ARIARunRequest):
    """Submit a goal and get ARIA's autonomous execution plan + result."""
    try:
        from orchestrator.aria_orchestrator import ARIAOrchestrator
        orchestrator = ARIAOrchestrator()
        result = orchestrator.run(req.goal)
        return {"status": "ok", "goal": req.goal, "result": result}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/aria/plan")
async def aria_plan(req: ARIAPlanRequest):
    """Submit an agent assignment during the planning phase."""
    return {
        "status": "ok",
        "agent_id": req.agent_id,
        "assigned_task": req.assigned_task,
        "priority": req.priority,
    }


@app.post("/aria/step")
async def aria_step(req: ARIAStepRequest):
    """Submit an oversight action during the execution phase."""
    return {
        "status": "ok",
        "action": req.action_type,
        "agent_id": req.agent_id,
    }


@app.get("/aria/state")
async def aria_state():
    return {"status": "ok", "system": "ARIA", "agents": 5, "phase": "ready"}


@app.get("/aria/report")
async def aria_report():
    from memory.episode_memory import get_stats
    return {"status": "ok", "memory_stats": get_stats()}


@app.post("/aria/evaluate")
async def aria_evaluate():
    return {"status": "ok", "message": "Evaluation complete"}


# ─────────────────────────────────────────────
# ARIA Multi-Agent RL Environment Routes
# ─────────────────────────────────────────────

_multi_env = None


class MultiEnvResetRequest(BaseModel):
    task_type: str = "synthesis_task"
    goal: str = ""
    seed: int = 42


class MultiEnvStepRequest(BaseModel):
    action_type: str = "monitor"
    worker_id: str = "worker_1"
    reason: Optional[str] = None


class MultiEnvPlanRequest(BaseModel):
    worker_id: str
    assigned_task_id: str
    priority: int = 3
    reason: Optional[str] = None


def _get_multi_env(task_type: str = "synthesis_task", seed: int = 42):
    global _multi_env
    if _multi_env is None:
        from rl_env import ARIAMultiEnv
        _multi_env = ARIAMultiEnv(task_type=task_type, seed=seed)
    return _multi_env


@app.post("/aria/multi/reset")
def multi_reset(req: MultiEnvResetRequest):
    """Reset the multi-agent RL environment for a new episode."""
    global _multi_env
    try:
        from rl_env import ARIAMultiEnv
        _multi_env = ARIAMultiEnv(task_type=req.task_type, seed=req.seed, goal=req.goal)
        obs = _multi_env.reset(goal=req.goal)
        return {
            "status": "ok",
            "task_type": req.task_type,
            "active_agents": [r.value for r in _multi_env.active_agents],
            "observation": obs.model_dump(),
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/aria/multi/plan")
def multi_plan(req: MultiEnvPlanRequest):
    """Submit a planning-phase task assignment."""
    env = _get_multi_env()
    try:
        from env.models import PlanningAction
        action = PlanningAction(
            worker_id=req.worker_id,
            assigned_task_id=req.assigned_task_id,
            priority=req.priority,
            reason=req.reason,
        )
        obs, reward, phase_done, info = env.plan(action)
        return {
            "status": "ok",
            "observation": obs.model_dump(),
            "reward": reward.model_dump(),
            "phase_done": phase_done,
            "message_bus": obs.message_bus.formatted(),
        }
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))


@app.post("/aria/multi/step")
def multi_step(req: MultiEnvStepRequest):
    """Submit an oversight action during the execution phase."""
    env = _get_multi_env()
    try:
        from env.models import OversightAction, OversightActionRequest
        action = OversightActionRequest(
            action_type=OversightAction(req.action_type),
            worker_id=req.worker_id,
            reason=req.reason,
        )
        obs, reward, done, info = env.step(action)
        return {
            "status": "ok",
            "observation": obs.model_dump(),
            "reward": reward.model_dump(),
            "reward_total": reward.total,
            "done": done,
            "message_bus": obs.message_bus.formatted(),
        }
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))


@app.get("/aria/multi/state")
def multi_state():
    """Current state of the multi-agent RL environment."""
    env = _get_multi_env()
    state = env.state()
    state["message_bus_view"] = env.message_bus.get_orchestrator_view().model_dump()
    return {"status": "ok", "state": state}


@app.get("/aria/multi/report")
def multi_report():
    """Full episode report including message bus statistics."""
    env = _get_multi_env()
    return {"status": "ok", "report": env.generate_run_report()}


@app.get("/aria/multi/tasks")
def multi_tasks():
    """List all available task types and their definitions."""
    from rl_env import TASK_REGISTRY, all_task_types
    result = {}
    for tt in all_task_types():
        td = TASK_REGISTRY[tt]
        result[tt.value] = {
            "fleet_task_id": td.fleet_task_id,
            "difficulty": td.difficulty,
            "active_agents": [r.value for r in td.active_agents],
            "n_anomalies": td.n_anomalies,
            "expected_steps": td.expected_steps,
            "description": td.description,
            "example_goal": td.example_goal,
        }
    return {"status": "ok", "tasks": result}


# ─────────────────────────────────────────────
# Legacy fleet routes (kept for compatibility)
# ─────────────────────────────────────────────

from env.oversight_env import FleetOversightEnv
from env.models import OversightAction, OversightActionRequest

_fleet_env: Optional[FleetOversightEnv] = None


def _get_fleet_env() -> FleetOversightEnv:
    global _fleet_env
    if _fleet_env is None:
        _fleet_env = FleetOversightEnv(task_id="easy_fleet", seed=42)
    return _fleet_env


class _FleetResetRequest(BaseModel):
    task_id: str = "easy_fleet"
    seed: int = 42


class _FleetStepRequest(BaseModel):
    action_type: str
    worker_id: str
    reason: Optional[str] = None


@app.post("/fleet/reset")
def fleet_reset(req: _FleetResetRequest):
    global _fleet_env
    _fleet_env = FleetOversightEnv(task_id=req.task_id, seed=req.seed)
    obs = _fleet_env.reset()
    return {"status": "ok", "observation": obs.model_dump(), "task_id": req.task_id}


@app.post("/fleet/step")
def fleet_step(req: _FleetStepRequest):
    env = _get_fleet_env()
    try:
        action = OversightActionRequest(
            action_type=OversightAction(req.action_type),
            worker_id=req.worker_id,
            reason=req.reason,
        )
    except ValueError as exc:
        raise HTTPException(status_code=400, detail=f"Invalid action: {exc}")
    obs, reward, done, info = env.step(action)
    return {
        "observation": obs.model_dump(),
        "reward": reward.model_dump(),
        "reward_total": reward.total,
        "done": done,
        "info": info,
    }


class FleetPlanRequest(BaseModel):
    worker_id: str
    assigned_task_id: str
    priority: int = 3
    reason: Optional[str] = None


@app.post("/fleet/plan")
def fleet_plan(req: FleetPlanRequest):
    env = _get_fleet_env()
    from env.models import PlanningAction
    action = PlanningAction(
        worker_id=req.worker_id,
        assigned_task_id=req.assigned_task_id,
        priority=req.priority,
        reason=req.reason,
    )
    obs, reward, phase_done, info = env.plan(action)
    return {
        "status": "ok",
        "observation": obs.model_dump(),
        "reward": reward.model_dump(),
        "reward_total": reward.total,
        "phase_done": phase_done,
        "info": info,
    }


@app.get("/fleet/state")
def fleet_state():
    return {"status": "ok", "state": _get_fleet_env().state()}


@app.get("/fleet/workers")
def fleet_workers():
    env = _get_fleet_env()
    workers_data = {}
    for wid in env.registry.workers:
        partial = env.registry.get_partial_obs(wid)
        workers_data[wid] = {
            **partial,
            "risk_score": env.registry.get_all_risk_scores().get(wid, 0.0),
            "is_anomalous": env.anomaly_injector.is_anomalous(wid),
        }
    return {"status": "ok", "workers": workers_data}


@app.get("/fleet/report")
def fleet_report():
    return {"status": "ok", "report": _get_fleet_env().generate_run_report()}


@app.post("/fleet/evaluate")
def fleet_evaluate():
    return {"status": "ok", "evaluation": _get_fleet_env().evaluate_run()}


# ─────────────────────────────────────────────
# Plot serving
# ─────────────────────────────────────────────

from fastapi.responses import FileResponse as _PlotFileResponse


@app.get("/plots/{filename}")
def serve_plot(filename: str):
    plot_path = _Path(__file__).resolve().parent / "plots" / filename
    if not plot_path.exists():
        raise HTTPException(status_code=404, detail=f"Plot {filename} not found.")
    return _PlotFileResponse(str(plot_path))


# ─────────────────────────────────────────────
# Gradio UI mount
# ─────────────────────────────────────────────

try:
    import gradio as gr
    from ui.gradio_demo import create_demo
    gradio_app = create_demo()
    app = gr.mount_gradio_app(app, gradio_app, path="/")
except Exception:
    @app.get("/", response_class=HTMLResponse)
    async def root():
        return HTMLResponse("<h1>ARIA — Autonomous Reasoning & Intelligence Architecture</h1><p>Start the app with <code>python app.py</code></p>")


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


if __name__ == "__main__":
    main()