import asyncio
import contextlib
import os
import secrets
import time
import uuid
import httpx
from contextlib import asynccontextmanager
from pathlib import Path
from typing import Annotated
from fastapi import FastAPI, HTTPException, Header, WebSocket, WebSocketDisconnect, Request
from fastapi.responses import FileResponse, Response, StreamingResponse
from fastapi.staticfiles import StaticFiles
from backend.engine import Engine
from backend.models import Event, Playback, Traffic, SMS, Dispatch
from backend.assets import handset_page
from backend.handset_access import public_origin, handset_redirect
from backend.teaching import authorize
from backend.recorder import Recorder, load_recording
from backend.csv_export import packet_snapshot, file_chunks, event_rows, event_csv

ROOT = Path(__file__).resolve().parents[1]
SCENARIOS = ["core-start", "ue-on", "pdu-session", "chat", "traffic-mix", "qos", "qos-reset", "auth-failure", "upf-stop", "upf-start", "deregister", "new-ue-register", "new-ue-session", "new-ue-data", "mobility-prepare", "mobility-xn", "mobility-n2"]

def create_app(replay=None, recording_dir=None):
    engine = Engine("replay" if replay else "live")
    controller = None
    scenario_lock = asyncio.Lock()
    from backend.factory import FactoryRuntime
    factory = FactoryRuntime(engine, lambda: get_controller())
    @asynccontextmanager
    async def lifespan(app):
        nonlocal controller
        if replay:
            engine.replay_events = load_recording(replay)
            engine.queue.extend(engine.replay_events)
            engine.replay_total = len(engine.replay_events)
        else:
            engine.recorder = Recorder(recording_dir or ROOT / "recordings")
        tasks = [asyncio.create_task(engine.pump()), asyncio.create_task(engine.tick())]
        if not replay:
            from backend.metrics import poll_metrics
            from backend.logs import tail_logs
            from backend.policy import enforce_policy
            tasks += [asyncio.create_task(poll_metrics(engine)), asyncio.create_task(tail_logs(engine)), asyncio.create_task(enforce_policy(engine))]
            tasks.append(asyncio.create_task(voice.observe()))
            if factory.enabled:
                tasks += [asyncio.create_task(factory.observe()), asyncio.create_task(factory.mobility())]
        try:
            yield
        finally:
            for task in tasks:
                task.cancel()
            await asyncio.gather(*tasks, return_exceptions=True)
            await factory.close()
            if engine.recorder:
                engine.recorder.close()
            if controller:
                controller.close()

    app = FastAPI(title="5G Core Classroom", lifespan=lifespan)
    secure_handset_origin = public_origin(os.getenv('PUBLIC_HANDSET_ORIGIN', ''))
    app.state.engine = engine
    app.state.factory = factory

    @app.get("/factory")
    async def factory_page():
        return FileResponse(ROOT / "dashboard/factory.html")

    @app.get("/api/factory/state")
    async def factory_state():
        return factory.status()

    @app.get("/api/factory/devices")
    async def factory_devices():
        return list(factory.registry.values())

    @app.post("/api/factory/control/{action}")
    async def factory_control(action: str):
        return await factory.control(action)

    @app.post("/api/factory/scenario/{name}")
    async def factory_scenario(name: str):
        return await factory.scenario(name)

    @app.post("/api/factory/speed")
    async def factory_speed(settings: dict):
        return await factory.speed(settings.get("rate"))

    @app.post("/api/factory/fault/{device}")
    async def factory_fault(device: str, settings: dict):
        return await factory.fault(device, settings)

    @app.post("/api/factory/handover/{device}/{mode}")
    async def factory_handover(device: str, mode: str):
        return await factory.handover(device, mode)

    @app.middleware("http")
    async def local_controls(request: Request, call_next):
        redirect = handset_redirect(request, secure_handset_origin)
        if redirect:
            return redirect
        # Prevent another website from making the browser operate the local core.
        if request.method in ("POST", "PUT", "PATCH", "DELETE") and request.url.path != "/api/ingest":
            origin = request.headers.get("origin")
            if origin and origin != f"{request.url.scheme}://{request.headers.get('host')}":
                from fastapi.responses import JSONResponse
                return JSONResponse({"detail": "Controls require the dashboard's own origin"}, status_code=403)
        response = await call_next(request)
        if request.url.path in ("/", "/mobile", "/voice") or request.url.path.startswith("/static/"):
            response.headers["Cache-Control"] = "no-cache"
        return response

    @app.get("/")
    async def index():
        return FileResponse(ROOT / "dashboard/index.html")

    @app.get("/presentation")
    async def presentation():
        return FileResponse(ROOT / "dashboard/presentation.html")

    @app.get("/mobile")
    async def mobile_page():
        return handset_page(ROOT, "mobile.html")

    @app.get("/source/{service}")
    async def source(service: str):
        # Only the lab's own application files; no arbitrary filesystem access.
        sources = {"app-server": "apps/server/__main__.py", "sms": "apps/sms.py",
                   "backend": "backend/main.py", "probe": "probe/decoders/packets.py",
                   "nef": "backend/nef_adapter.py", "af": "backend/demo_af.py", "observer": "backend/control.py", "ran": "mobility/ran_agent.py", "ran-patch": "mobility/patch_ueransim.py"}
        if service not in sources:
            raise HTTPException(404, "Unknown application source")
        return FileResponse(ROOT / sources[service], media_type="text/plain; charset=utf-8")

    @app.get("/api/health")
    async def health():
        return {"status": "ok", "mode": engine.mode}

    @app.get("/api/state")
    async def state():
        return engine.snapshot()

    def recording_source():
        return Path(replay) if replay else engine.recorder.path if engine.recorder else None

    @app.get("/api/export/status")
    async def export_status():
        source = recording_source()
        packet = packet_snapshot(recording_dir or ROOT / "recordings") if engine.mode == "live" else None
        return {"mode": engine.mode, "events_available": source is not None and source.is_file(),
                "recording": source.name if source else None, "packets_available": packet is not None,
                "packets": packet[2] if packet else None}

    @app.get("/api/export/events.csv")
    async def export_events():
        source = recording_source()
        if source is None or not source.is_file():
            raise HTTPException(503, "Event recording is not available")
        # Ingest writes complete lines synchronously on this event loop. Bound
        # the live snapshot now; later appends and playback changes cannot affect it.
        limit = source.stat().st_size
        rows = (e.model_dump() for e in engine.replay_events) if replay else event_rows(source, limit)
        return StreamingResponse(event_csv(rows), media_type="text/csv; charset=utf-8",
            headers={"Content-Disposition": f'attachment; filename="metriczone-events-{source.stem}.csv"',
                     "Cache-Control": "no-store"})

    @app.get("/api/export/packets.csv")
    async def export_packets():
        if engine.mode != "live":
            raise HTTPException(409, "Replay exports decoded events; individual frames require a live packet recording")
        packet = packet_snapshot(recording_dir or ROOT / "recordings")
        if packet is None:
            raise HTTPException(503, "Packet recording is not available; start the updated capture probe")
        path, limit, metadata = packet
        return StreamingResponse(file_chunks(path, limit), media_type="text/csv; charset=utf-8",
            headers={"Content-Disposition": f'attachment; filename="metriczone-{path.name}"',
                     "Content-Length": str(limit), "X-Packet-Rows": str(metadata["rows"]), "Cache-Control": "no-store"})

    @app.get("/metrics")
    async def metrics():
        lines = ["# HELP demo_n3_packets_total GTP-U data packets observed on the core bridge", "# TYPE demo_n3_packets_total counter"]
        lines += [f'demo_n3_packets_total{{direction="{direction}"}} {count}' for direction, count in engine.captured_packets.items()]
        lines += ["# HELP demo_n3_bytes_total Inner IP bytes observed in GTP-U", "# TYPE demo_n3_bytes_total counter"]
        lines += [f'demo_n3_bytes_total{{direction="{direction}"}} {count}' for direction, count in engine.captured_bytes.items()]
        return Response("\n".join(lines) + "\n", media_type="text/plain; version=0.0.4")

    @app.get("/api/scenarios")
    async def scenarios():
        return {"scenarios": SCENARIOS, "enabled": engine.mode == "live"}

    @app.post("/api/ingest")
    async def ingest(events: list[Event], token: Annotated[str | None, Header(alias="X-Ingest-Token")] = None):
        if engine.mode != "live":
            raise HTTPException(409, "Capture ingestion is disabled during replay")
        if not token or not secrets.compare_digest(token, os.getenv("INGEST_TOKEN", "local-demo-capture")):
            raise HTTPException(403, "Invalid capture token")
        if len(events) > 1000:
            raise HTTPException(413, "Maximum 1000 events per batch")
        if any(event.protocol not in ("NEF", "FACTORY") for event in events):
            engine.probe_last_seen = time.time()
        events = sorted(events, key=lambda e: e.ts)
        if engine.recorder:
            engine.recorder.write_batch(events)
        for event in events:
            voice.capture(event)
            engine.ingest(event, record=False)
        return {"accepted": len(events)}

    @app.post("/api/playback")
    async def playback(settings: Playback):
        engine.set_playback(settings.speed, settings.paused)
        return engine.snapshot()

    @app.post("/api/step")
    async def step():
        engine.set_playback(engine.speed, True)
        engine.steps += 1
        engine.wake.set()
        return {"stepping": True}

    @app.post("/api/replay/restart")
    async def restart():
        if engine.mode != "replay":
            raise HTTPException(409, "Launch backend with --replay first")
        engine.paused = True
        await engine.restart_replay()
        return engine.snapshot()

    def get_controller():
        nonlocal controller
        if engine.mode != "live":
            raise HTTPException(409, "Docker controls are disabled during replay")
        if not controller:
            from backend.control import Controller
            controller = Controller()
        return controller

    @app.post("/scenario/{name}")
    async def scenario(name: str):
        if name not in SCENARIOS:
            raise HTTPException(404, "Unknown scenario")
        if scenario_lock.locked():
            raise HTTPException(409, "Another scenario is running")
        async with scenario_lock:
            try:
                if name in ("mobility-xn", "mobility-n2"):
                    mode = name.removeprefix("mobility-")
                    control = get_controller()
                    await asyncio.to_thread(control.mobility_enabled)
                    engine.ingest(Event(ts=time.time(), src="ue6", dst="observer", plane="control",
                        protocol="MOBILITY", name=f"{mode.upper()} handover requested", ue="imsi-001010000000006",
                        procedure="handover", raw={"phase": "measuring", "mode": mode, "source": "controller"},
                        explain="Demo control: begin real tunnel-bound UDP echoes and one TCP connection, then trigger native handover. This row is not a network signalling packet."))
                    result = await asyncio.to_thread(control.handover, mode, dict(engine.live_radio.get("imsi-001010000000006", {})))
                    expected = {"Xn Handover Request", "Xn Handover Request Acknowledge", "Path Switch Request", "Path Switch Request Acknowledge", "Xn UE Context Release"} if mode == "xn" else {
                        "Handover Required", "Handover Request", "Handover Request Acknowledge", "Handover Command", "Handover Notify", "UE Context Release Command", "UE Context Release Complete"}
                    # Capture is independent of playback speed/pause. Give its normal
                    # batch delivery time to arrive before checking protocol evidence.
                    await asyncio.sleep(0.7)
                    captured = {e.name for e in engine.radio_evidence if e.ue == "imsi-001010000000006" and
                                result["started_at"] - .05 <= e.ts <= result["completed_at"]}
                    core_expected = {"Update SM Context", "Update SM Context Response", "PFCP Session Modification Request", "PFCP Session Modification Response"}
                    core_captured = {e.name for e in engine.radio_evidence if e.name in core_expected and
                                     result["started_at"] - .05 <= e.ts <= result["completed_at"]}
                    result["captured_signalling"] = sorted(expected & captured)
                    result["captured_core_path_update"] = sorted(core_captured)
                    result["missing_signalling"] = sorted((expected - captured) | (core_expected - core_captured))
                    result["unexpected_reestablishment"] = sorted(captured & {"Registration Request", "PDU Session Establishment Request", "PFCP Session Establishment Request"})
                    if result["missing_signalling"]:
                        result["success"] = False
                        result["output"] = result["output"].replace(": PASS", ": FAIL", 1) + " · missing packet evidence: " + ", ".join(result["missing_signalling"])
                    if result["unexpected_reestablishment"]:
                        result["success"] = False
                        result["output"] = result["output"].replace(": PASS", ": FAIL", 1) + " · unexpected registration/session re-establishment"
                    result["phase"] = "complete" if result["success"] else "failed"
                    result["source"] = "measurement"
                    import json
                    artifact = ROOT / "recordings" / result["evidence_file"]
                    evidence = json.loads(artifact.read_text(encoding="utf-8"))
                    evidence["result"] = result
                    artifact.write_text(json.dumps(evidence, indent=2), encoding="utf-8")
                    engine.ingest(Event(ts=time.time(), src="ue6", dst="app-server", plane="user", protocol="MOBILITY",
                        name=f"{mode.upper()} continuity {'PASS' if result['success'] else 'FAIL'}", ue="imsi-001010000000006",
                        procedure="handover", raw=result, explain="Measured UDP loss and echo gaps, plus preservation of the UE IP, TUN interface and one TCP connection. This summary is measurement evidence, not a generated control-plane packet."))
                else:
                    result = await asyncio.to_thread(get_controller().scenario, name)
                return {"scenario": name, **result}
            except HTTPException:
                raise
            except Exception as exc:
                raise HTTPException(503, str(exc)) from exc

    @app.get("/api/mobility/evidence/{name}")
    async def mobility_evidence(name: str):
        import re
        if not re.fullmatch(r"mobility-[a-f0-9]{32}\.json", name) and name not in ("mobility-xn-measurement.json", "mobility-n2-measurement.json"):
            raise HTTPException(404, "Unknown measurement")
        path = ROOT / "recordings" / name
        if not path.is_file():
            raise HTTPException(404, "Unknown measurement")
        return FileResponse(path, media_type="application/json", filename=name)

    @app.post("/api/traffic")
    async def traffic(request: Traffic):
        if scenario_lock.locked():
            raise HTTPException(409, "Another scenario is running")
        async with scenario_lock:
            try:
                return await asyncio.to_thread(get_controller().traffic, request)
            except HTTPException:
                raise
            except Exception as exc:
                raise HTTPException(503, str(exc)) from exc

    @app.post("/api/sms/connect")
    async def connect_phones():
        if scenario_lock.locked():
            raise HTTPException(409, "Another scenario is running")
        async with scenario_lock:
            try:
                return await asyncio.to_thread(get_controller().connect_phones)
            except HTTPException:
                raise
            except Exception as exc:
                raise HTTPException(503, str(exc)) from exc

    @app.get("/api/sms")
    async def messages():
        return {"messages": list(engine.sms_messages), "mode": engine.mode}

    @app.post("/api/sms")
    async def sms(request: SMS):
        message = {"id": uuid.uuid4().hex, **request.model_dump(), "sent_at": time.time()}
        return await send_sms(request, message)

    async def send_sms(request, message, metadata=None):
        if scenario_lock.locked():
            raise HTTPException(409, "Another scenario is running")
        async with scenario_lock:
            control = get_controller()
            try:
                receipt = await asyncio.to_thread(control.sms, request, message)
            except Exception as exc:
                failed = {**message, **(metadata or {}), "status": "failed", "error": str(exc), "source": "ue-app"}
                engine.ingest(Event(ts=time.time(), src=request.sender, dst=request.recipient, plane="user",
                    protocol="SMS", name="SMS Failed", procedure="data", raw=failed,
                    explain="The sending phone did not receive a delivery receipt."))
                raise HTTPException(503, str(exc)) from exc
            receipt["source"] = "ue-receipt"
            receipt.update(metadata or {})
            engine.ingest(Event(ts=receipt["delivered_at"], src=request.sender, dst=request.recipient,
                plane="user", protocol="SMS", name="SMS Delivered", procedure="data", raw=receipt,
                explain=f"{request.recipient.upper()} received the short text from {request.sender.upper()} through their 5G data tunnels."))
            return receipt

    async def af_request(method, path, payload=None):
        if engine.mode != "live":
            raise HTTPException(409, "NEF demo controls are disabled during replay")
        async with httpx.AsyncClient() as client:
            try:
                response = await client.request(method, os.getenv("AF_URL", "http://172.24.0.10:8000") + path,
                    json=payload, headers={"X-Teaching-Token": os.getenv("TEACHING_TOKEN", "local-nef-teaching")}, timeout=20)
                if response.is_error:
                    raise HTTPException(response.status_code, response.json().get("detail", response.text))
                return response.json()
            except (httpx.HTTPError, ValueError) as exc:
                raise HTTPException(503, f"Metriczone AF is unavailable: {exc}") from exc

    @app.get("/api/nef")
    async def nef_status():
        if engine.mode == "replay":
            return {"mode": "replay", "teaching": True, "demo": engine.nef_demo, "subscriptions": []}
        try:
            demo = await af_request("GET", "/teaching/v1/demo")
            async with httpx.AsyncClient() as client:
                response = await client.get(os.getenv("NEF_URL", "http://172.24.0.2:8000") + "/teaching/v1/subscriptions",
                    headers={"X-AF-Key": os.getenv("NEF_AF_TOKEN", "local-metriczone-af")}, timeout=5)
                response.raise_for_status()
                return {"mode": "live", "teaching": True, "demo": demo, "subscriptions": response.json(), "healthy": True}
        except (HTTPException, httpx.HTTPError, ValueError) as exc:
            return {"mode": "live", "teaching": True, "healthy": False, "error": str(getattr(exc, "detail", exc))}

    @app.post("/api/nef/prepare")
    async def nef_prepare(request: SMS):
        if scenario_lock.locked():
            raise HTTPException(409, "Another scenario is running")
        async with scenario_lock:
            await af_request("POST", "/teaching/v1/reset")
            try:
                result = await asyncio.to_thread(get_controller().prepare_meter, request)
            except HTTPException:
                raise
            except Exception as exc:
                raise HTTPException(503, str(exc)) from exc
            engine.ingest(Event(ts=time.time(), src="observer", dst=request.recipient, plane="control", protocol="NEF",
                name="Prepare offline meter", procedure="exposure", raw={"teaching": True, "demo": result},
                explain="The lab operator stops the target UE container. Subscribe and queue, then bring the phone online."))
            return result

    @app.post("/api/nef/subscribe")
    async def nef_subscribe(request: SMS):
        if scenario_lock.locked():
            raise HTTPException(409, "Another scenario is running")
        async with scenario_lock:
            return await af_request("POST", "/teaching/v1/demo", request.model_dump())

    @app.post("/api/nef/online")
    async def nef_online(request: SMS):
        if scenario_lock.locked():
            raise HTTPException(409, "Another scenario is running")
        async with scenario_lock:
            demo = await af_request("GET", "/teaching/v1/demo")
            if demo.get("phase") not in ("queued", "notified") or any(demo.get(key) != getattr(request,key) for key in ("sender", "recipient", "text")):
                raise HTTPException(409, "Subscribe and queue this command before bringing its phone online")
            try:
                return await asyncio.to_thread(get_controller().connect_pair, request)
            except Exception as exc:
                raise HTTPException(503, str(exc)) from exc

    @app.post("/api/nef/reset")
    async def nef_reset():
        if scenario_lock.locked():
            raise HTTPException(409, "Another scenario is running")
        async with scenario_lock:
            return await af_request("POST", "/teaching/v1/reset")

    @app.get("/api/nef/observe")
    async def nef_observe(sender: str, recipient: str, x_teaching_token: str | None = Header(default=None)):
        authorize(x_teaching_token)
        if sender not in ("ue1", "ue2", "ue3", "ue4") or recipient not in ("ue1", "ue2", "ue3", "ue4") or sender == recipient:
            raise HTTPException(422, "Choose two distinct working phones")
        if scenario_lock.locked():
            raise HTTPException(409, "A phone operation is in progress; availability is unknown")
        try:
            # Read-only probes must not reserve the controls and reject ordinary SMS.
            observation = await asyncio.to_thread(get_controller().observe, sender, recipient)
            if scenario_lock.locked():
                raise HTTPException(409, "A phone operation overlapped the probe; availability is unknown")
            return observation
        except HTTPException:
            raise
        except Exception as exc:
            raise HTTPException(503, f"Availability observation failed: {exc}") from exc

    @app.post("/api/nef/dispatch")
    async def nef_dispatch(request: Dispatch, x_teaching_token: str | None = Header(default=None)):
        authorize(x_teaching_token)
        message = request.model_dump(exclude={"notification_id", "subscription_id"})
        sms_request = SMS(sender=request.sender, recipient=request.recipient, text=request.text)
        return await send_sms(sms_request, message, {"trigger": "nef", "command_id": request.id,
            "notification_id": request.notification_id, "subscription_id": request.subscription_id})

    from backend.mobile import MobileFleet
    fleet = MobileFleet(engine, get_controller, scenario_lock, send_sms)
    app.state.mobile = fleet.default
    app.state.mobile_fleet = fleet
    app.include_router(fleet.router)
    from backend.voice import Voice
    voice = Voice(engine, get_controller, scenario_lock, ROOT)
    voice.fleet = fleet
    app.state.voice = voice
    app.include_router(voice.router)

    @app.websocket("/ws")
    async def websocket(ws: WebSocket):
        origin = ws.headers.get("origin")
        if origin and origin.split("://", 1)[-1] != ws.headers.get("host"):
            await ws.close(code=1008)
            return
        await ws.accept()
        queue = asyncio.Queue(maxsize=1500)
        engine.clients.add(queue)
        async def sender():
            while True:
                payload = await queue.get()
                if payload["type"] == "disconnect":
                    await ws.close(code=1013)
                    return
                await ws.send_json(payload)
        async def receiver():
            while True:
                await ws.receive_text()
        try:
            events = {event["id"]: event for event in (*engine.sequence_history, *engine.history)}
            await ws.send_json({"type": "hello", "state": engine.snapshot(),
                               "events": sorted(events.values(), key=lambda event: event["id"])})
            tasks = [asyncio.create_task(sender()), asyncio.create_task(receiver())]
            await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
        except (WebSocketDisconnect, RuntimeError, asyncio.CancelledError):
            pass
        finally:
            engine.clients.discard(queue)
            for task in locals().get("tasks", []):
                task.cancel()
            with contextlib.suppress(asyncio.CancelledError):
                await asyncio.gather(*locals().get("tasks", []), return_exceptions=True)

    app.mount("/static", StaticFiles(directory=ROOT / "dashboard"), name="static")
    return app