#!/usr/bin/env python3
"""
main.py — Cloud Run entrypoint, OSVauco OPAX.
Modes: light (gemini-2.5-flash) | heavy (gemini-2.5-pro)
ML-1: telemetri, state store, DAG.
CG1+CG2: billing endpoints + IAP-beskyttelse.
CG3: static-mappe serveres fra /static/*.
AUTH: /auth/login + /auth/callback for klient OAuth2 onboarding.
CG5a: POST /notify/webhook — push til Slack/Teams/Chat/Discord/etc.
CG5b: POST /notify/email  — SendGrid digest til kunde-e-post.
"""

import os
import sys
import time
import pathlib

sys.path.insert(0, os.path.join(os.path.dirname(__file__), "agents", "core-logic"))

from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import FileResponse, JSONResponse, RedirectResponse, Response
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel, Field
from typing import List, Optional, Dict, Any
import datetime
import httpx
from cachetools import cached, TTLCache

from agent import run, authorize_mode
from ml import build_agent_dag, execute_dag, get_store, log_agent_call
from ml.telemetry import log_dag_execution
from ml.billing_agent import BillingAgent
from agents.aws_billing_agent import AWSBillingAgent
from ml.anomaly_detector import AnomalyDetector
from auth.token_store import save_token


from google.cloud import firestore
import firebase_admin
from firebase_admin import credentials, messaging
import base64
import json
from functools import wraps
from starlette.middleware.sessions import SessionMiddleware
from uvicorn.middleware.proxy_headers import ProxyHeadersMiddleware
from authlib.integrations.starlette_client import OAuth
import sendgrid
from sendgrid.helpers.mail import Mail

AGENT_ID = "osvauco-opax"
app = FastAPI(
    title="OSVauco OPAX Agent",
    description="Agent for OSVauco-OPAX platform.",
    version="0.1.0",
)

# Paths exempt from IAP enforcement (health/readiness probes reach Cloud Run
# directly without the IAP-injected x-goog-authenticated-user-email header).
IAP_EXEMPT_PATHS = {"/health", "/healthz", "/readiness", "/liveness"}

@app.middleware("http")
async def require_iap(request: Request, call_next):
    if request.url.path in IAP_EXEMPT_PATHS:
        return await call_next(request)
    if not request.headers.get("x-goog-authenticated-user-email"):
        return Response(status_code=401, content="Unauthorized")
    return await call_next(request)

# ── FIREBASE & FIRESTORE INIT ───────────────────────────────────────────────
db = None
try:
    firebase_admin.initialize_app()
    db = firestore.Client()
    print("Firestore client initialized successfully.")
except Exception as e:
    print(f"WARNING: Firestore client failed to initialize: {e}", file=sys.stderr)

# ── AUTH & SESSION ────────────────────────────────────────────────────────────
app.add_middleware(ProxyHeadersMiddleware, trusted_hosts="*")
app.add_middleware(SessionMiddleware, secret_key=os.environ.get("SESSION_SECRET"))

oauth = OAuth()
oauth.register(
    name='google',
    client_id=os.environ.get("GOOGLE_CLIENT_ID"),
    client_secret=os.environ.get("GOOGLE_CLIENT_SECRET"),
    server_metadata_url='https://accounts.google.com/.well-known/openid-configuration',
    client_kwargs={'scope': 'openid email profile'}
)


@app.get('/auth/login')
async def login(request: Request):
    redirect_uri = "https://opax.vauco.no/auth/callback"
    return await oauth.google.authorize_redirect(request, redirect_uri)

@app.get('/auth/callback', name='auth')
async def auth(request: Request):
    token = await oauth.google.authorize_access_token(request)
    user = token.get('userinfo')
    if user:
        request.session['user'] = dict(user)
    return RedirectResponse(url='/static/billing-dashboard.html')


    @app.get("/auth/me")
    async def me(request: Request):
        user = request.session.get("user")
        if not user:
            raise HTTPException(status_code=401, detail="Not authenticated")
        return JSONResponse(user)


    @app.get('/auth/logout')
    async def logout(request: Request):
        request.session.pop('user', None)
        return RedirectResponse(url='/static/billing-dashboard.html')


ALLOWED_EMAILS = [email.strip() for email in os.environ.get("ALLOWED_EMAILS", "").split(",") if email.strip()]
ALERT_EMAIL = os.environ.get("ALERT_EMAIL")


def require_auth(func):
    """Decorator to protect endpoints that require authentication."""
    @wraps(func)
    async def wrapper(request: Request, *args, **kwargs):
        user = request.session.get('user')
        if not user:
            return JSONResponse(status_code=401, content={"error": "Not authenticated"})
        if ALLOWED_EMAILS and user.get('email') not in ALLOWED_EMAILS:
            return JSONResponse(status_code=403, content={"error": "Email not allowed"})
        return await func(request, *args, **kwargs)
    return wrapper


# ── STATIC FILES ──────────────────────────────────────────────────────────────
_static_dir = pathlib.Path(__file__).parent / "static"
if _static_dir.is_dir():
    app.mount("/static", StaticFiles(directory=str(_static_dir), html=True), name="static")


# ── MODELLER ──────────────────────────────────────────────────────────────────
class RunRequest(BaseModel):
    message: str
    user_id: str = "opax"
    session_id: str = "default"
    mode: str = "light"


class DagRequest(BaseModel):
    messages: List[str] = Field(...)
    user_id: str = "opax"
    session_id: str = "default"
    mode: str = "light"
    scheduler: str = Field("threads")

class PushSubscription(BaseModel):
    token: str
    budget_nok: float

class BudgetWebhookPayload(BaseModel):
    message: dict
    subscription: str

class BudgetUpdateRequest(BaseModel):
    budget: float


# ── NOTIFICATION MODELS (CG5a + CG5b) ────────────────────────────────────────
class WebhookNotifyRequest(BaseModel):
    url: str = Field(..., description="Webhook-URL (Slack/Teams/Google Chat/Discord/custom)")
    event: str = Field("custom", description="Event-type: spike | digest | budget | anomaly | custom")
    title: str = Field("CostGuard varsel", description="Tittel på meldingen")
    body: str = Field(..., description="Meldingstekst")
    payload: Optional[Dict[str, Any]] = Field(None, description="Valgfri rådata å inkludere")

class EmailNotifyRequest(BaseModel):
    to: Optional[str] = Field(None, description="Mottaker-e-post. Faller tilbake til ALERT_EMAIL hvis tom.")
    subject: Optional[str] = Field(None, description="Emne. Auto-generert fra event hvis tom.")
    event: str = Field("digest", description="Event-type: spike | digest | budget | anomaly | custom")
    body_html: Optional[str] = Field(None, description="HTML-innhold. Auto-generert fra billing-data hvis tom.")
    payload: Optional[Dict[str, Any]] = Field(None, description="Valgfri rådata å inkludere i e-post")


# ── NOTIFICATION SERVICE (CG5a + CG5b) ───────────────────────────────────────
class NotificationService:
    """
    Felles delivery layer for CostGuard-varsler.
    CG5a: send_webhook()  — poster til hvilken som helst webhook-URL
    CG5b: send_email()    — sender via SendGrid
    Brukt av POST /notify/webhook og POST /notify/email.
    """

    # Event-type → emoji + standardtekst
    EVENT_META = {
        "spike":   {"emoji": "🚨", "label": "Kostnadsspike oppdaget"},
        "digest":  {"emoji": "📊", "label": "Daglig kostnadsoppsummering"},
        "budget":  {"emoji": "⚠️",  "label": "Budsjettgrense nærmer seg"},
        "anomaly": {"emoji": "🔍", "label": "Anomali oppdaget"},
        "custom":  {"emoji": "📢", "label": "CostGuard varsel"},
    }

    def _event_meta(self, event: str) -> dict:
        return self.EVENT_META.get(event, self.EVENT_META["custom"])

    def _build_slack_payload(self, req: WebhookNotifyRequest) -> dict:
        """Slack/Teams/Google Chat-kompatibelt JSON-format."""
        meta = self._event_meta(req.event)
        return {
            "text": f"{meta['emoji']} *{req.title}*\n{req.body}",
            "attachments": [
                {
                    "color": "#FF6B35" if req.event in ("spike", "budget") else "#4A90D9",
                    "fields": [
                        {"title": k, "value": str(v), "short": True}
                        for k, v in (req.payload or {}).items()
                    ]
                }
            ] if req.payload else []
        }

    async def send_webhook(self, req: WebhookNotifyRequest) -> dict:
        """
        CG5a: POST JSON til oppgitt webhook-URL.
        Støtter Slack, Teams, Google Chat, Discord og alle custom webhooks.
        """
        slack_payload = self._build_slack_payload(req)
        try:
            async with httpx.AsyncClient(timeout=10.0) as client:
                resp = await client.post(
                    req.url,
                    json=slack_payload,
                    headers={"Content-Type": "application/json"},
                )
            print(f"[NotificationService] webhook → {req.url} status={resp.status_code}")
            if resp.status_code >= 400:
                return {"status": "error", "http_status": resp.status_code, "detail": resp.text[:200]}
            return {"status": "ok", "http_status": resp.status_code, "event": req.event}
        except httpx.TimeoutException:
            return {"status": "error", "detail": "Webhook timeout (>10s)"}
        except Exception as e:
            return {"status": "error", "detail": str(e)}

    async def send_email(self, req: EmailNotifyRequest) -> dict:
        """
        CG5b: Send e-post via SendGrid.
        Hvis body_html er None, bygges innhold fra billing-data automatisk.
        """
        api_key = os.environ.get("SENDGRID_API_KEY")
        to_email = req.to or ALERT_EMAIL

        if not api_key:
            return {"status": "not_configured", "detail": "SENDGRID_API_KEY mangler"}
        if not to_email:
            return {"status": "not_configured", "detail": "Ingen mottaker-e-post (to eller ALERT_EMAIL)"}

        meta = self._event_meta(req.event)
        subject = req.subject or f"{meta['emoji']} CostGuard: {meta['label']} — {datetime.date.today()}"

        # Auto-generer HTML-innhold fra billing-data hvis ikke oppgitt
        if req.body_html:
            html_content = req.body_html
        else:
            html_content = await self._build_digest_html(req.payload)

        try:
            message = Mail(
                from_email="costguard@osvauco.no",
                to_emails=to_email,
                subject=subject,
                html_content=html_content,
            )
            sg = sendgrid.SendGridAPIClient(api_key)
            response = sg.send(message)
            msg_id = response.headers.get("X-Message-Id", "")
            print(f"[NotificationService] email → {to_email} msg_id={msg_id}")
            return {"status": "ok", "message_id": msg_id, "to": to_email, "event": req.event}
        except Exception as e:
            print(f"[NotificationService] email error: {e}", file=sys.stderr)
            return {"status": "error", "detail": str(e)}

    async def _build_digest_html(self, extra_payload: Optional[dict] = None) -> str:
        """Henter billing-data og bygger HTML-digest."""
        try:
            billing = BillingAgent()
            summary = billing.get_summary().get("summary", [])
            forecast = billing.get_forecast()
            mtd = forecast.get("month_to_date_cost", 0)
            daily_avg = forecast.get("daily_average_last_7_days", 0)
            top3 = "".join(
                f"<li>{s['service']}: kr {s.get('total_cost', 0):.2f}</li>"
                for s in summary[:3]
            )
        except Exception:
            mtd, daily_avg, top3 = 0, 0, "<li>Data ikke tilgjengelig</li>"

        extra_rows = ""
        if extra_payload:
            extra_rows = "".join(
                f"<tr><td style='padding:4px 8px;color:#666'>{k}</td>"
                f"<td style='padding:4px 8px'><strong>{v}</strong></td></tr>"
                for k, v in extra_payload.items()
            )

        return f"""
        <div style='font-family:sans-serif;max-width:600px;margin:0 auto'>
          <div style='background:#1a1a2e;padding:20px;border-radius:8px 8px 0 0'>
            <h2 style='color:#fff;margin:0'>📊 CostGuard — Daglig oppsummering</h2>
            <p style='color:#aaa;margin:4px 0 0'>{datetime.date.today()}</p>
          </div>
          <div style='background:#f9f9f9;padding:20px;border-radius:0 0 8px 8px'>
            <table style='width:100%;border-collapse:collapse'>
              <tr><td style='padding:4px 8px;color:#666'>Kostnad MTD</td>
                  <td style='padding:4px 8px'><strong>kr {mtd:.2f}</strong></td></tr>
              <tr><td style='padding:4px 8px;color:#666'>Daglig snitt (7d)</td>
                  <td style='padding:4px 8px'><strong>kr {daily_avg:.2f}</strong></td></tr>
              {extra_rows}
            </table>
            <h4 style='margin-top:16px'>Topp 3 GCP-tjenester</h4>
            <ul>{top3}</ul>
            <hr style='border:none;border-top:1px solid #ddd;margin:16px 0'>
            <p style='color:#999;font-size:12px'>CostGuard by Vauco AS &middot; opax.vauco.no</p>
          </div>
        </div>
        """


# Singleton
_notifier = NotificationService()


# ── NOTIFY ENDPOINTS (CG5a + CG5b) ───────────────────────────────────────────

async def _notify_webhook(request: Request, req: WebhookNotifyRequest):
    result = await _notifier.send_webhook(req)
    status_code = 200 if result["status"] == "ok" else 502
    return JSONResponse(status_code=status_code, content=result)

app.add_api_route("/notify/webhook", endpoint=require_auth(_notify_webhook), methods=["POST"])


async def _notify_email(request: Request, req: EmailNotifyRequest):
    result = await _notifier.send_email(req)
    if result["status"] == "not_configured":
        return JSONResponse(status_code=503, content=result)
    status_code = 200 if result["status"] == "ok" else 500
    return JSONResponse(status_code=status_code, content=result)

app.add_api_route("/notify/email", endpoint=require_auth(_notify_email), methods=["POST"])


# ── ROOT LANDING PAGE ───────────────────────────────────────────────────
@app.get("/")
def root():
    return FileResponse("static/opax.html")


# ── HEALTH ────────────────────────────────────────────────────────────────────
@app.get("/health")
def health():
    return {"status": "ok"}


@app.get("/manifest.json", include_in_schema=False)
def manifest():
    return FileResponse("static/manifest.json")


@app.get("/sw.js", include_in_schema=False)
def service_worker():
    return FileResponse("static/sw.js")


# ── AUTH — User Authentication ──────────────────────────────────────────
@app.get('/admin')
@require_auth
async def admin_panel(request: Request):
    return FileResponse("static/admin.html")

@app.post('/admin/create-customer')
@require_auth
async def create_customer(request: Request):
    import subprocess, shlex
    data = await request.json()
    customer_name = data.get("customer_name", "").strip()
    project_id = data.get("project_id", "").strip()
    billing_account_id = data.get("billing_account_id", "").strip()
    alert_email = data.get("alert_email", "").strip()
    container_image = data.get("container_image", "").strip()
    region = data.get("region", "europe-north1").strip()

    if not all([customer_name, project_id, billing_account_id, alert_email, container_image]):
        raise HTTPException(status_code=400, detail="Alle felt er påkrevd")

    import re
    if not re.match(r'^[a-z0-9_-]+$', customer_name):
        raise HTTPException(status_code=400, detail="customer_name kan kun inneholde a-z, 0-9, - og _")

    customer_dir = f"infrastructure/terraform/customers/{customer_name}"
    template_dir = "infrastructure/terraform/customers/_template"

    import os, shutil
    if os.path.exists(customer_dir):
        raise HTTPException(status_code=409, detail=f"Kunde {customer_name} eksisterer allerede")

    shutil.copytree(template_dir, customer_dir)

    tfvars_content = f'''customer_id           = "{customer_name}"
project_id            = "{project_id}"
region                = "{region}"
billing_account_id    = "{billing_account_id}"
alert_email           = "{alert_email}"
billing_viewer_emails = ["chris.christiansen@vauco.no", "jason.vauger@vauco.no"]
container_image       = "{container_image}"
'''
    with open(f"{customer_dir}/terraform.tfvars", "w") as f:
        f.write(tfvars_content)

    try:
        init = subprocess.run(
            ["terraform", f"-chdir={customer_dir}", "init", "-no-color"],
            capture_output=True, text=True, timeout=120
        )
        if init.returncode != 0:
            shutil.rmtree(customer_dir)
            raise HTTPException(status_code=500, detail=f"terraform init feilet: {init.stderr[-500:]}")

        apply = subprocess.run(
            ["terraform", f"-chdir={customer_dir}", "apply", "-auto-approve", "-no-color"],
            capture_output=True, text=True, timeout=600
        )
        if apply.returncode != 0:
            raise HTTPException(status_code=500, detail=f"terraform apply feilet: {apply.stderr[-500:]}")

        return {"status": "ok", "customer": customer_name, "project_id": project_id}

    except subprocess.TimeoutExpired:
        raise HTTPException(status_code=504, detail="Terraform tok for lang tid (>10 min)")


# ── BILLING ENDPOINTS (CG1 + CG2) ────────────────────────────────────────────

async def get_budget(request: Request):
    if not db:
        return JSONResponse(status_code=500, content={"error": "Firestore is not configured"})
    try:
        doc_ref = db.collection("settings").document("budget")
        doc = doc_ref.get()
        if doc.exists:
            return JSONResponse(content={"budget": doc.to_dict().get("limit", 500)})
        else:
            return JSONResponse(content={"budget": 500})
    except Exception as e:
        return JSONResponse(status_code=500, content={"error": str(e)})

async def set_budget(request: Request, payload: BudgetUpdateRequest):
    if not db:
        return JSONResponse(status_code=500, content={"error": "Firestore is not configured"})
    try:
        doc_ref = db.collection("settings").document("budget")
        doc_ref.set({"limit": payload.budget})
        return JSONResponse(content={"status": "ok", "budget": payload.budget})
    except Exception as e:
        return JSONResponse(status_code=500, content={"error": str(e)})

app.add_api_route("/billing/budget", endpoint=require_auth(get_budget), methods=["GET"])
app.add_api_route("/billing/budget", endpoint=require_auth(set_budget), methods=["POST"])


@app.post("/billing/email-report")
async def trigger_email_report():
    """
    Legacy scheduler-endepunkt. Kaller NotificationService.send_email() internt.
    Beholdes for bakoverkompatibilitet med Cloud Scheduler.
    """
    result = await _notifier.send_email(EmailNotifyRequest(event="digest"))
    if result["status"] == "not_configured":
        return JSONResponse(status_code=200, content={"status": "not_configured", "error": result["detail"]})
    if result["status"] == "error":
        return JSONResponse(status_code=500, content=result)
    return JSONResponse(content=result)


@app.post("/billing/snapshot")
async def create_daily_snapshot(request: Request):
    email_response = await trigger_email_report()
    if email_response.status_code != 200:
        print(f"Daily snapshot: email report failed. Status: {email_response.status_code}", file=sys.stderr)
    return JSONResponse(content={"status": "ok", "snapshot_id": str(datetime.date.today())})


async def get_history(request: Request):
    if not db:
        return JSONResponse(status_code=500, content={"error": "Firestore is not configured"})
    try:
        end_date = datetime.date.today()
        start_date = end_date - datetime.timedelta(days=90)
        docs = db.collection("daily_snapshots") \
            .where("created_at", ">=", start_date.isoformat()) \
            .order_by("created_at", direction=firestore.Query.DESCENDING) \
            .limit(90) \
            .stream()
        history = [{"date": doc.id, **doc.to_dict()} for doc in docs]
        history.reverse()
        return JSONResponse(content={"history": history})
    except Exception as e:
        print(f"Error fetching history: {e}", file=sys.stderr)
        return JSONResponse(status_code=500, content={"error": str(e)})

app.add_api_route("/billing/history", endpoint=require_auth(get_history), methods=["GET"])


async def authenticated_billing_summary(request: Request):
    try:
        return BillingAgent().get_summary()
    except Exception as exc:
        return JSONResponse(status_code=500, content={"error": str(exc)})
app.add_api_route("/billing/summary", endpoint=require_auth(authenticated_billing_summary), methods=["GET"])


async def authenticated_billing_forecast(request: Request):
    try:
        return BillingAgent().get_forecast()
    except Exception as exc:
        return JSONResponse(status_code=500, content={"error": str(exc)})
app.add_api_route("/billing/forecast", endpoint=require_auth(authenticated_billing_forecast), methods=["GET"])


async def authenticated_billing_anomalies(request: Request):
    try:
        return AnomalyDetector().detect_anomalies()
    except Exception as exc:
        return JSONResponse(status_code=500, content={"error": str(exc)})
app.add_api_route("/billing/anomalies", endpoint=require_auth(authenticated_billing_anomalies), methods=["GET"])


@app.get("/billing-dashboard")
def billing_dashboard_view(request: Request):
    return RedirectResponse(url='/static/billing-dashboard.html')

@app.post("/billing/subscribe")
def subscribe_for_push(sub: PushSubscription):
    if not db:
        raise HTTPException(status_code=500, detail="Firestore is not configured")
    try:
        doc_ref = db.collection("push_subscribers").document(sub.token)
        doc_ref.set({"budget_nok": sub.budget_nok, "subscribed_at": firestore.SERVER_TIMESTAMP})
        return {"status": "ok"}
    except Exception as e:
        print(f"Error subscribing token: {e}", file=sys.stderr)
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/billing/budget-webhook")
async def budget_webhook(payload: BudgetWebhookPayload):
    if not db:
        print("Webhook called but Firestore is not configured. Aborting.", file=sys.stderr)
        return {"status": "error", "detail": "Firestore not configured"}
    try:
        data = base64.b64decode(payload.message.get("data", "")).decode("utf-8")
        data_json = json.loads(data)
        cost_amount = data_json.get("costAmount", 0)
        budget_amount = data_json.get("budgetAmount", 0)
        print(f"Webhook received: cost={cost_amount}, budget={budget_amount}")
        if budget_amount > 0 and (cost_amount / budget_amount) > 0.8:
            percent_used = round((cost_amount / budget_amount) * 100)
            subscribers_ref = db.collection("push_subscribers")
            tokens = [subscriber.id for subscriber in subscribers_ref.stream()]
            if not tokens:
                print("Budget threshold exceeded, but no push subscribers found.")
                return {"status": "no subscribers"}
            notification = messaging.Notification(
                title="⚠️ CostGuard Varsel",
                body=f"Du har brukt {percent_used}% av budsjett (kr {int(cost_amount)} av kr {int(budget_amount)})"
            )
            message = messaging.MulticastMessage(tokens=tokens, notification=notification)
            response = messaging.send_multicast(message)
            print(f"Sent push notification to {response.success_count} subscribers.")
        else:
            print("Budget threshold not exceeded, no notification sent.")
        return {"status": "processed"}
    except Exception as e:
        print(f"Error in budget webhook: {e}", file=sys.stderr)
        return {"status": "error", "detail": str(e)}


@app.get("/billing/live")
@cached(TTLCache(maxsize=1, ttl=3600))
async def billing_live(request: Request):
    user = request.session.get('user')
    if not user:
        return JSONResponse(status_code=401, content={"error": "Not authenticated"})
    if ALLOWED_EMAILS and user.get('email') not in ALLOWED_EMAILS:
        return JSONResponse(status_code=403, content={"error": "Email not allowed"})
    print("Billing Budget API not implemented, falling back to BigQuery", file=sys.stderr)
    bq_forecast = BillingAgent().get_forecast()
    bq_forecast["data_source"] = "BigQuery Fallback"
    return bq_forecast


# ── AWS BILLING ENDPOINTS ─────────────────────────────────────────────────────
async def authenticated_aws_billing_summary(request: Request):
    try:
        return AWSBillingAgent().get_summary()
    except Exception as exc:
        return JSONResponse(status_code=500, content={"error": str(exc)})
app.add_api_route("/billing/aws/summary", endpoint=require_auth(authenticated_aws_billing_summary), methods=["GET"])

async def authenticated_aws_billing_forecast(request: Request):
    try:
        return AWSBillingAgent().get_forecast()
    except Exception as exc:
        return JSONResponse(status_code=500, content={"error": str(exc)})
app.add_api_route("/billing/aws/forecast", endpoint=require_auth(authenticated_aws_billing_forecast), methods=["GET"])


# ── AGENT ENDPOINTS ───────────────────────────────────────────────────────────
@app.post("/run")
def run_agent(req: RunRequest):
    try:
        authorize_mode(req.user_id, req.mode)
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    except PermissionError as e:
        raise HTTPException(status_code=403, detail=str(e))

    store = get_store()
    start = time.monotonic()
    error_msg = None
    response = None

    try:
        response = run(
            message=req.message,
            user_id=req.user_id,
            session_id=req.session_id,
            mode=req.mode,
        )
    except Exception as e:
        error_msg = str(e)
        raise HTTPException(status_code=500, detail=error_msg)
    finally:
        duration = round(time.monotonic() - start, 3)
        success = error_msg is None
        model = "gemini-2.5-flash" if req.mode in ("light", "A") else "gemini-2.5-pro"
        log_agent_call(
            agent_id=AGENT_ID,
            input_payload={"message": req.message, "mode": req.mode},
            output=response,
            model_used=model,
            mode=req.mode,
            duration_s=duration,
            success=success,
            error=error_msg,
        )
        store.push(AGENT_ID, "last_duration_s", duration)
        store.push(AGENT_ID, "last_mode", req.mode)
        store.push(AGENT_ID, "last_success", success)
        if response:
            store.push(AGENT_ID, "last_result_preview", response[:200])

    return {"response": response}


@app.post("/run/dag")
def run_dag(req: DagRequest):
    try:
        authorize_mode(req.user_id, req.mode)
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    except PermissionError as e:
        raise HTTPException(status_code=403, detail=str(e))

    def make_agent_fn(msg: str, idx: int):
        def _agent_fn(payload: dict):
            return run(
                message=msg,
                user_id=payload["user_id"],
                session_id=f"{payload['session_id']}-dag-{idx}",
                mode=payload["mode"],
            )
        _agent_fn.__name__ = f"opax-dag-{idx}"
        return _agent_fn

    payload = {"user_id": req.user_id, "session_id": req.session_id, "mode": req.mode}
    agent_fns = [make_agent_fn(msg, i) for i, msg in enumerate(req.messages)]
    agent_ids = [f"opax-dag-{i}" for i in range(len(req.messages))]

    start = time.monotonic()
    tasks = build_agent_dag(agent_fns, payload, agent_ids)
    results = execute_dag(tasks, scheduler=req.scheduler)
    total_dur = round(time.monotonic() - start, 3)

    store = get_store()
    log_dag_execution(dag_id=req.session_id, agent_results=results, total_duration_s=total_dur)
    store.push(AGENT_ID, "last_dag_total_duration_s", total_dur)
    store.push(AGENT_ID, "last_dag_agent_count", len(agents))
    store.push(AGENT_ID, "last_dag_success_count", sum(1 for r in results if r["success"]))

    return {
        "total_duration_s": total_dur,
        "results": [
            {
                "index": i,
                "message": req.messages[i],
                "response": r["result"],
                "success": r["success"],
                "duration_s": r["duration_s"],
                "error": r["error"],
            }
            for i, r in enumerate(results)
        ],
    }


@app.get("/state")
def get_state():
    return get_store().snapshot()

@app.get("/state/agents")
def list_agents():
    return {"agents": get_store().list_agents()}

@app.get("/state/aggregate/{key}")
def aggregate_key(key: str):
    return get_store().aggregate(key)

@app.get("/telemetry/history")
def telemetry_history(limit: int = 50):
    return {"history": get_store().history(limit=limit)}


if __name__ == "__main__":
    import uvicorn
    port = int(os.environ.get("PORT", 8080))
    uvicorn.run(app, host="0.0.0.0", port=port)