#!/usr/bin/env python3
"""
main.py — Cloud Run entrypoint for OSVauco OPAX agent.
Exposes a FastAPI HTTP API that wraps agent.run().

ML-1 changes:
  - /run now records timing + outcome via ml.telemetry
  - /run pushes last_result + last_duration_s to AgentStateStore
  - /run/dag accepts a list of messages and runs them as a parallel Dask DAG
  - /state and /telemetry endpoints expose live ML observability
"""

import os
import sys
import time

# Ensure agents/core-logic is on the path
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "agents", "core-logic"))

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field
from typing import List, Optional

from agent import run, authorize_mode  # agents/core-logic/agent.py
from ml import (
    build_agent_dag,
    execute_dag,
    get_store,
    log_agent_call,
)
from ml.telemetry import log_dag_execution

app = FastAPI(title="OSVauco OPAX Agent")

AGENT_ID = "opax-core"


# ---------------------------------------------------------------------------
# Request / Response models
# ---------------------------------------------------------------------------

class RunRequest(BaseModel):
    message: str
    user_id: str = "opax"
    session_id: str = "default"
    mode: str = "A"


class DagRequest(BaseModel):
    messages: List[str] = Field(..., description="Liste av meldinger som kjøres parallelt")
    user_id: str = "opax"
    session_id: str = "default"
    mode: str = "A"
    scheduler: str = Field("threads", description="Dask scheduler: synchronous | threads | processes")


# ---------------------------------------------------------------------------
# Health
# ---------------------------------------------------------------------------

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


# ---------------------------------------------------------------------------
# /run — enkelt agent-kall med telemetri + state store (ML-1)
# ---------------------------------------------------------------------------

@app.post("/run")
def run_agent(req: RunRequest):
    authorize_mode(req.user_id, req.mode)  # raises PermissionError if unauthorized

    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 PermissionError as e:
        raise HTTPException(status_code=403, detail=str(e))
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    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

        # ML-1: logg til JSONL-telemetri
        log_agent_call(
            agent_id=AGENT_ID,
            input_payload={"message": req.message, "mode": req.mode},
            output=response,
            model_used="gemini-2.0-flash" if req.mode == "A" else "gemini-2.5-pro",
            mode=req.mode,
            duration_s=duration,
            success=success,
            error=error_msg,
        )

        # ML-1: push til Parameter Server state store
        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}


# ---------------------------------------------------------------------------
# /run/dag — parallell Dask DAG over flere meldinger (ML-1)
# ---------------------------------------------------------------------------

@app.post("/run/dag")
def run_dag(req: DagRequest):
    """
    Kjører en liste av meldinger som parallelle Dask-tasks.
    Hver melding behandles av en separat agent-wrapper.
    Returnerer alle svar med timing og suksess-status.
    """
    try:
        authorize_mode(req.user_id, req.mode)
    except PermissionError as e:
        raise HTTPException(status_code=403, detail=str(e))

    def make_agent_fn(msg: str, idx: int):
        """Lukker over msg + idx for å lage en unik callable per melding."""
        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)

    # ML-1: logg DAG-kjøring og push aggregert state
    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(results))
    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)
        ],
    }


# ---------------------------------------------------------------------------
# /state — live Parameter Server state snapshot (ML-1 observability)
# ---------------------------------------------------------------------------

@app.get("/state")
def get_state():
    """Returnerer nåværende AgentStateStore snapshot."""
    return get_store().snapshot()


@app.get("/state/agents")
def list_agents():
    """Returnerer alle agent-IDer som har pushet state."""
    return {"agents": get_store().list_agents()}


@app.get("/state/aggregate/{key}")
def aggregate_key(key: str):
    """Aggregerer alle agenter sine verdier for en gitt nøkkel."""
    return get_store().aggregate(key)


# ---------------------------------------------------------------------------
# /telemetry — historikk-log (ML-1 observability)
# ---------------------------------------------------------------------------

@app.get("/telemetry/history")
def telemetry_history(limit: int = 50):
    """Returnerer siste N state-operasjoner fra AgentStateStore historikk."""
    return {"history": get_store().history(limit=limit)}


# ---------------------------------------------------------------------------
# Startup
# ---------------------------------------------------------------------------

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