"""
Vitalis Unified Telemetry System (VUTS) Server — Vitalis FSI
FastAPI Asynchronous Backend supporting state polling and real-time WebSocket tracing.
"""
from __future__ import annotations
import json
import asyncio
from pathlib import Path
from typing import Set
from fastapi import FastAPI, Request, WebSocket, WebSocketDisconnect
from fastapi.responses import HTMLResponse
from fastapi.templating import Jinja2Templates

BASE_DIR = Path(__file__).resolve().parent.parent.parent
TEMPLATES_DIR = Path(__file__).resolve().parent / "templates"

app = FastAPI(title="Vitalis Core Telemetry System")
templates = Jinja2Templates(directory=TEMPLATES_DIR)

class TelemetryBroker:
    def __init__(self):
        self.active_connections: Set[WebSocket] = set()

    async def connect(self, websocket: WebSocket):
        await websocket.accept()
        self.active_connections.add(websocket)

    def disconnect(self, websocket: WebSocket):
        self.active_connections.remove(websocket)

    async def broadcast(self, message: dict):
        if not self.active_connections:
            return
        payload = json.dumps(message)
        disconnected = set()
        for connection in self.active_connections:
            try:
                await connection.send_text(payload)
            except Exception:
                disconnected.add(connection)
        
        for conn in disconnected:
            self.active_connections.remove(conn)

broker = TelemetryBroker()

@app.get("/", response_class=HTMLResponse)
async def get_dashboard(request: Request):
    """Pass valid FastAPI Request object context down to the Jinja2 template instance."""
    return templates.TemplateResponse("index.html", {"request": request})

@app.get("/api/state")
async def get_state():
    """Fallback REST endpoint to pull static diagnostic records from the filesystem."""
    self_model_path = BASE_DIR / "self_model.json"
    memory_path = BASE_DIR / "lorein_episodic_memory.json"
    
    state = {"self_model": {}, "memory_chain": []}
    
    if self_model_path.exists():
        with open(self_model_path, "r", encoding="utf-8") as f:
            state["self_model"] = json.load(f)
            
    if memory_path.exists():
        with open(memory_path, "r", encoding="utf-8") as f:
            try:
                state["memory_chain"] = json.load(f)[-10:]
            except Exception:
                pass
                
    return state

@app.websocket("/ws/telemetry")
async def websocket_endpoint(websocket: WebSocket):
    await broker.connect(websocket)
    try:
        state = await get_state()
        await websocket.send_text(json.dumps({"event": "SYSTEM_SYNC", "data": state}))
        while True:
            await websocket.receive_text()
    except WebSocketDisconnect:
        broker.disconnect(websocket)