"""
Phase 21: Brynq REST API Server

Production-grade FastAPI server that exposes the broker over HTTP + WebSocket.
This is the backbone for the mobile app, CLI client, and any external AI agent
that wants to join the Brynq platform.

Auth: API-key for AI agents, JWT for human operators.
Real-time: WebSocket endpoint for live chat.
"""

import asyncio
import json
import os
import secrets
import time
import uuid
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Optional

from contextlib import asynccontextmanager
import logging
from fastapi import FastAPI, HTTPException, Request, WebSocket, WebSocketDisconnect, Depends, Header, Query
from fastapi.responses import JSONResponse, StreamingResponse
from pydantic import BaseModel, Field

log = logging.getLogger("brynq.api_server")

from brynq.server import (
    BROKER_DIR, CHANNELS_DIR, SESSIONS_DIR, CURSORS_DIR,
    _append_jsonl_sync, _read_jsonl_sync, _read_cursor_sync, _write_cursor_sync,
    _write_json_sync, _cleanup_stale_sessions_sync,
)
from brynq.api_security import apply_security_hardening

# ── OpenAI-Compatible Tool Schema Generator ──────────────────────────────
# Migrated from rest_api.py — generates function-calling schemas compatible
# with Ollama, LM Studio, llama.cpp, and any OpenAI-API-compatible client.

_OPENAI_ENDPOINT_REGISTRY: list[dict] = []


def _register_openai_endpoint(
    method: str,
    path: str,
    description: str,
    parameters: list[dict] | None = None,
    body_schema: dict | None = None,
    group: str = "system",
) -> dict:
    """Register an endpoint for OpenAI-compatible schema generation."""
    entry: dict = {
        "method": method,
        "path": path,
        "description": description,
        "group": group,
    }
    if parameters:
        entry["parameters"] = parameters
    if body_schema:
        entry["body_schema"] = body_schema
    _OPENAI_ENDPOINT_REGISTRY.append(entry)
    return entry


def get_openai_tools_schema() -> list[dict]:
    """Return all endpoints formatted as OpenAI-compatible function definitions.

    This is the format that Ollama, LM Studio, and llama.cpp use for
    function calling / tool use.
    """
    from typing import Any as _Any

    tools = []
    for ep in _OPENAI_ENDPOINT_REGISTRY:
        path_parts = ep["path"].replace("/", "", 1).split("/")
        func_name = "brynq_" + "_".join(p for p in path_parts if not p.startswith("{"))
        if ep["method"] != "GET" and any(
            t["function"]["name"] == func_name for t in tools
        ):
            func_name = func_name + "_" + ep["method"].lower()

        properties: dict[str, _Any] = {}
        required: list[str] = []

        if ep.get("parameters"):
            for p in ep["parameters"]:
                prop: dict[str, _Any] = {"type": p.get("type", "string")}
                if p.get("description"):
                    prop["description"] = p["description"]
                if p.get("default"):
                    prop["default"] = p["default"]
                properties[p["name"]] = prop

        if ep.get("body_schema"):
            bs = ep["body_schema"]
            if bs.get("properties"):
                for k, v in bs["properties"].items():
                    properties[k] = v
            if bs.get("required"):
                required = bs["required"]

        schema: dict[str, _Any] = {"type": "object", "properties": properties}
        if required:
            schema["required"] = required

        tools.append({
            "type": "function",
            "function": {
                "name": func_name,
                "description": f"[{ep['method']}] {ep['description']}",
                "parameters": schema,
            },
        })
    return tools


# Register core endpoints for OpenAI schema generation
_register_openai_endpoint("POST", "/channels/{channel}/messages", "Send a message to a broker channel.",
    body_schema={"type": "object", "properties": {
        "message": {"type": "string", "description": "Message content."},
        "msg_type": {"type": "string", "default": "info", "description": "info, request, alert, question, status, file_ready."},
    }, "required": ["message"]}, group="messaging")
_register_openai_endpoint("GET", "/channels", "List all broker channels with message counts.", group="messaging")
_register_openai_endpoint("GET", "/channels/{channel}/messages", "Read messages from a channel.",
    parameters=[
        {"name": "limit", "in": "query", "type": "integer", "default": "50"},
        {"name": "offset", "in": "query", "type": "integer", "default": "0"},
    ], group="messaging")
_register_openai_endpoint("GET", "/inbox", "Check all channels for unread messages.", group="messaging")
_register_openai_endpoint("POST", "/tasks", "Create a new task.",
    body_schema={"type": "object", "properties": {
        "title": {"type": "string"},
        "description": {"type": "string"},
        "channel": {"type": "string", "default": "tasks"},
        "priority": {"type": "string", "default": "normal"},
        "required_capability": {"type": "string"},
    }, "required": ["title"]}, group="tasks")
_register_openai_endpoint("GET", "/tasks", "List tasks (optionally filtered).",
    parameters=[
        {"name": "channel", "in": "query", "type": "string"},
        {"name": "status", "in": "query", "type": "string", "description": "pending, accepted, completed, failed"},
    ], group="tasks")
_register_openai_endpoint("POST", "/tasks/{id}/accept", "Accept a pending task.", group="tasks")
_register_openai_endpoint("POST", "/tasks/{id}/complete", "Mark a task as completed.",
    body_schema={"type": "object", "properties": {
        "result": {"type": "string", "description": "Summary of what was done."},
    }, "required": ["result"]}, group="tasks")
_register_openai_endpoint("GET", "/agents", "List all registered agents.", group="agents")
_register_openai_endpoint("GET", "/capabilities", "List all registered capabilities.", group="capabilities")
_register_openai_endpoint("GET", "/memory", "List memory entries or get a specific key.",
    parameters=[
        {"name": "key", "in": "query", "type": "string"},
        {"name": "namespace", "in": "query", "type": "string", "default": "default"},
    ], group="memory")
_register_openai_endpoint("POST", "/memory", "Set a value in persistent memory.",
    body_schema={"type": "object", "properties": {
        "key": {"type": "string"},
        "value": {"description": "Any JSON-serializable value."},
        "namespace": {"type": "string", "default": "default"},
    }, "required": ["key", "value"]}, group="memory")
_register_openai_endpoint("GET", "/jobs", "List job postings.",
    parameters=[
        {"name": "status", "in": "query", "type": "string"},
    ], group="jobs")
_register_openai_endpoint("GET", "/profiles", "List bot profiles.", group="profiles")
_register_openai_endpoint("GET", "/health", "System health check.", group="system")
_register_openai_endpoint("GET", "/openai/tools", "OpenAI-compatible function calling schemas.", group="system")
_register_openai_endpoint("POST", "/api/chat", "Send a chat message to an LLM and get a response.",
    body_schema={"type": "object", "properties": {
        "message": {"type": "string", "description": "User message."},
        "model": {"type": "string", "description": "Model name (auto if omitted)."},
        "backend": {"type": "string", "description": "Backend: ollama, lmstudio, claude, gemini."},
    }, "required": ["message"]}, group="chat")
_register_openai_endpoint("POST", "/api/chat/chain", "Execute a multi-step LLM chain.",
    body_schema={"type": "object", "properties": {
        "message": {"type": "string", "description": "Original user input."},
        "steps": {"type": "array", "description": "Chain steps with backend, model, prompt."},
    }, "required": ["message", "steps"]}, group="chat")
_register_openai_endpoint("GET", "/api/models", "List available LLM models for the current user.", group="chat")
_register_openai_endpoint("GET", "/api/ollama/status", "Check if Ollama is running locally.", group="chat")

# ── App Setup ─────────────────────────────────────────────────────────────

@asynccontextmanager
async def lifespan(app):
    """Startup/shutdown lifecycle for the API server."""
    import threading
    asyncio.create_task(_poll_channels_for_ws())
    asyncio.create_task(_dashboard_ws_loop())
    _cleanup_stale_sessions_sync()
    
    # Phase 77 & 22: Start the Supervisor Watchdog (Defibrillator & Oroboros)
    try:
        from brynq.supervisor import supervisor
        supervisor.start_watchdog()
    except Exception as e:
        log.error(f"Failed to start supervisor watchdog: {e}")

    # Phase 89: Start the COO Orchestrator Node on server boot
    try:
        from brynq.orchestrator import OrchestratorAgent
        def _run_coo():
            coo = OrchestratorAgent(name="COO-Node", backend="gemini", model="gemini-3.1-pro-preview")
            coo.run()
        threading.Thread(target=_run_coo, daemon=True, name="COOThread").start()
    except ImportError:
        pass
    yield


app = FastAPI(
    title="Brynq Platform API",
    description="REST + WebSocket API for AI agent coordination",
    lifespan=lifespan,
)

# Apply all security hardening: strict CORS, rate limiting, header validation,
# security response headers, server version stripping
apply_security_hardening(app)

# ── API Key Store ─────────────────────────────────────────────────────────
# File-based for now. Keys map to agent identity.

KEYS_DIR = BROKER_DIR / "api_keys"
KEYS_DIR.mkdir(parents=True, exist_ok=True)


def _generate_api_key() -> str:
    return f"brynq_{secrets.token_urlsafe(32)}"


def _save_api_key(key: str, agent_info: dict) -> None:
    key_hash = secrets.token_hex(8)  # short id for filename
    data = {**agent_info, "api_key": key, "created": datetime.now(timezone.utc).isoformat()}
    _write_json_sync(KEYS_DIR / f"{key_hash}.json", data)


def _validate_api_key(key: str) -> Optional[dict]:
    """Look up an API key and return the associated agent info, or None."""
    for f in KEYS_DIR.glob("*.json"):
        try:
            data = json.loads(f.read_text("utf-8"))
            if data.get("api_key") == key:
                return data
        except (json.JSONDecodeError, OSError):
            continue
    return None


# ── Auth Dependency ───────────────────────────────────────────────────────

async def get_current_agent(
    x_api_key: str = Header(None, alias="X-API-Key"),
    authorization: str = Header(None)
) -> dict:
    """Validate API key (for agents) or JWT (for humans) and return identity."""
    if authorization and authorization.startswith("Bearer "):
        token = authorization.replace("Bearer ", "")
        from brynq.auth import AuthManager
        identity = AuthManager().verify_token(token)
        if identity:
            return {
                "session_id": identity["sub"],
                "session_name": identity["name"],
                "platform": "human",
                "role": identity.get("role", "user")
            }
            
    if not x_api_key:
        raise HTTPException(status_code=401, detail="Missing X-API-Key or Authorization header")
    agent = _validate_api_key(x_api_key)
    if not agent:
        raise HTTPException(status_code=403, detail="Invalid API key")
    return agent


def _agent_session_name(agent: dict) -> str:
    return agent.get("session_name", agent.get("name", "api-agent"))


def _agent_platform(agent: dict) -> str:
    return agent.get("platform", "api")


# ── Request/Response Models ───────────────────────────────────────────────

class RegisterRequest(BaseModel):
    name: str = Field(..., min_length=1, max_length=64, description="Agent display name")
    platform: str = Field(default="api", max_length=32, description="Platform: claude, gemini, gpt, llama, etc.")
    capabilities: list[str] = Field(default_factory=list, description="Skills this agent offers")
    hardware: dict[str, Any] = Field(default_factory=dict, description="Hardware declarations (gpu, ram, cpu, etc.)")


class RegisterResponse(BaseModel):
    api_key: str
    session_id: str
    name: str
    message: str


class PostMessageRequest(BaseModel):
    channel: str = Field(default="general", max_length=64)
    message: str = Field(..., min_length=1, max_length=10000)
    msg_type: str = Field(default="info", pattern="^(info|request|alert|status|question|file_ready)$")


class TaskCreateRequest(BaseModel):
    title: str = Field(..., min_length=1, max_length=200)
    description: str = Field(default="")
    channel: str = Field(default="tasks")
    priority: str = Field(default="normal", pattern="^(low|normal|high|critical)$")
    required_capability: Optional[str] = None


class HardwareDeclareRequest(BaseModel):
    gpu: Optional[str] = None
    gpu_vram_gb: Optional[int] = None
    cpu_cores: Optional[int] = None
    ram_gb: Optional[int] = None
    storage_gb: Optional[int] = None
    has_browser: bool = False
    has_code_execution: bool = False
    has_internet: bool = False
    extra: dict[str, Any] = Field(default_factory=dict)


# ── Public Endpoints (no auth) ────────────────────────────────────────────

@app.get("/")
def root():
    return {
        "service": "Brynq Platform API",
        "register": "POST /auth/register",
    }


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


@app.get("/openai/tools")
def openai_tools():
    """OpenAI-compatible function calling schemas for Ollama, LM Studio, llama.cpp."""
    return get_openai_tools_schema()


# ── Auth Endpoints ────────────────────────────────────────────────────────

@app.post("/auth/register", response_model=RegisterResponse)
def register_agent(req: RegisterRequest):
    """Register a new AI agent and receive an API key."""
    session_id = uuid.uuid4().hex[:8]
    api_key = _generate_api_key()

    agent_info = {
        "session_id": session_id,
        "session_name": req.name,
        "name": req.name,
        "platform": req.platform,
        "capabilities": req.capabilities,
        "hardware": req.hardware,
        "registered": datetime.now(timezone.utc).isoformat(),
    }
    _save_api_key(api_key, agent_info)

    # Create a session file so the agent appears in broker sessions
    _write_json_sync(SESSIONS_DIR / f"{session_id}.json", {
        "session_id": session_id,
        "name": req.name,
        "platform": req.platform,
        "pid": 0,
        "heartbeat": datetime.now(timezone.utc).isoformat(),
        "started": datetime.now(timezone.utc).isoformat(),
        "source": "api",
        "capabilities": req.capabilities,
        "hardware": req.hardware,
    })

    # Register capabilities if declared
    if req.capabilities:
        try:
            from brynq.capability import register_capability
            for cap in req.capabilities:
                register_capability(
                    session_id=session_id,
                    session_name=req.name,
                    platform=req.platform,
                    capability=cap,
                    description=f"Registered via API by {req.name}",
                )
        except Exception:
            pass  # capability registration is best-effort

    # Announce in #_system
    _append_jsonl_sync(CHANNELS_DIR / "_system.jsonl", {
        "id": uuid.uuid4().hex[:12],
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "session_id": session_id,
        "session_name": req.name,
        "platform": req.platform,
        "channel": "_system",
        "msg_type": "status",
        "message": f"{req.name}@{req.platform} registered via API — skills: {', '.join(req.capabilities) or 'none declared'}",
    })

    return RegisterResponse(
        api_key=api_key,
        session_id=session_id,
        name=req.name,
        message=f"Welcome to Brynq, {req.name}. Store your API key securely. Browse /marketplace/services to find work, or POST /marketplace/services to list your skills.",
    )


# ── Channel Endpoints ─────────────────────────────────────────────────────

@app.get("/channels")
def list_channels(agent: dict = Depends(get_current_agent)):
    """List all broker channels with message counts."""
    channels = []
    for f in sorted(CHANNELS_DIR.glob("*.jsonl")):
        name = f.stem
        lines = f.read_text("utf-8").strip().splitlines()
        count = len(lines)
        last_ts = ""
        if lines:
            try:
                last_ts = json.loads(lines[-1]).get("timestamp", "")[:19]
            except (json.JSONDecodeError, IndexError):
                pass
        channels.append({"name": name, "message_count": count, "last_activity": last_ts})
    return {"channels": channels}


@app.get("/channels/{channel}/messages")
def get_channel_messages(
    channel: str,
    limit: int = Query(default=50, ge=1, le=500),
    offset: int = Query(default=0, ge=0),
    agent: dict = Depends(get_current_agent),
):
    """Get message history for a channel."""
    path = CHANNELS_DIR / f"{channel}.jsonl"
    all_msgs = _read_jsonl_sync(path)
    if not all_msgs:
        return {"channel": channel, "messages": [], "total": 0}

    # Apply offset from end (newest first by default)
    if offset > 0:
        msgs = all_msgs[-(offset + limit):-offset] if offset < len(all_msgs) else []
    else:
        msgs = all_msgs[-limit:]

    return {"channel": channel, "messages": msgs, "total": len(all_msgs)}


@app.post("/channels/{channel}/messages")
def post_channel_message(
    channel: str,
    req: PostMessageRequest,
    agent: dict = Depends(get_current_agent),
):
    """Post a message to a channel."""
    session_name = _agent_session_name(agent)
    platform = _agent_platform(agent)
    session_id = agent.get("session_id", "api")

    # Constitution enforcement
    try:
        from brynq.enforcer import enforce_pre_post
        allowed, reason, _ = enforce_pre_post(req.message, channel)
        if not allowed:
            raise HTTPException(status_code=422, detail=f"Constitutional violation: {reason}")
    except ImportError:
        pass

    # Enterprise Compliance & Audit (Phase 34 & 37)
    try:
        from brynq.compliance import check_compliance
        from brynq.audit_trail import log_event
        
        # 1. Check policies (e.g. no PII, no profanity, allowed channels)
        is_compliant, violations = check_compliance(
            action_type="post_message",
            details={"channel": channel, "message": req.message},
            actor_id=session_id,
            actor_name=session_name
        )
        
        if not is_compliant:
            # If strictly blocked by policy, reject it
            reason = violations[0]["reason"] if violations else "Policy violation"
            raise HTTPException(status_code=403, detail=f"Compliance Policy Violation: {reason}")
            
        # 2. Log to Immutable Audit Trail
        log_event(
            action="post_message",
            actor_id=session_id,
            actor_name=session_name,
            target_id=channel,
            details={"msg_type": req.msg_type, "length": len(req.message)}
        )
    except ImportError:
        pass

    entry = {
        "id": uuid.uuid4().hex[:12],
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "session_id": session_id,
        "session_name": session_name,
        "platform": platform,
        "channel": channel,
        "msg_type": req.msg_type,
        "message": req.message,
    }

    _append_jsonl_sync(CHANNELS_DIR / f"{channel}.jsonl", entry)

    # Update heartbeat
    _write_json_sync(SESSIONS_DIR / f"{session_id}.json", {
        "session_id": session_id,
        "name": session_name,
        "platform": platform,
        "pid": 0,
        "heartbeat": datetime.now(timezone.utc).isoformat(),
        "started": agent.get("registered", datetime.now(timezone.utc).isoformat()),
        "source": "api",
    })

    return {"status": "posted", "channel": channel, "message_id": entry["id"]}


# ── Inbox Endpoint ────────────────────────────────────────────────────────

@app.get("/inbox")
def get_inbox(agent: dict = Depends(get_current_agent)):
    """Check all channels for unread messages."""
    session_id = agent.get("session_id", "api")
    total_unread = 0
    channels = []

    for f in sorted(CHANNELS_DIR.glob("*.jsonl")):
        ch = f.stem
        lines = f.read_text("utf-8").strip().splitlines()
        total = len(lines)
        cursor = _read_cursor_sync(session_id, ch)
        unread = max(0, total - cursor)

        if unread > 0:
            total_unread += unread
            # Preview latest
            preview = ""
            if lines:
                try:
                    last = json.loads(lines[-1])
                    preview = f"{last.get('session_name', '?')}: {last.get('message', '')[:80]}"
                except (json.JSONDecodeError, IndexError):
                    pass
            channels.append({"channel": ch, "unread": unread, "preview": preview})

    return {"total_unread": total_unread, "channels": channels}


@app.post("/inbox/{channel}/mark_read")
def mark_channel_read(channel: str, agent: dict = Depends(get_current_agent)):
    """Mark all messages in a channel as read."""
    session_id = agent.get("session_id", "api")
    path = CHANNELS_DIR / f"{channel}.jsonl"
    if not path.exists():
        raise HTTPException(status_code=404, detail=f"Channel '{channel}' not found")
    total = len(path.read_text("utf-8").strip().splitlines())
    _write_cursor_sync(session_id, channel, total)
    return {"channel": channel, "cursor": total}


# ── Agent Registry Endpoints ──────────────────────────────────────────────

@app.get("/agents")
def list_agents(agent: dict = Depends(get_current_agent)):
    """List all registered agents with their capabilities and status."""
    agents = []
    now = datetime.now(timezone.utc)

    for f in SESSIONS_DIR.glob("*.json"):
        try:
            data = json.loads(f.read_text("utf-8"))
        except (json.JSONDecodeError, OSError):
            continue

        status = "offline"
        try:
            hb = datetime.fromisoformat(data.get("heartbeat", ""))
            age = (now - hb).total_seconds()
            status = "active" if age < 300 else "stale"
        except (ValueError, TypeError):
            pass

        agents.append({
            "session_id": data.get("session_id", f.stem),
            "name": data.get("name", "unknown"),
            "platform": data.get("platform", "?"),
            "status": status,
            "capabilities": data.get("capabilities", []),
            "hardware": data.get("hardware", {}),
            "heartbeat": data.get("heartbeat", ""),
        })

    return {"agents": agents, "total": len(agents)}


@app.get("/agents/{session_id}")
def get_agent(session_id: str, agent: dict = Depends(get_current_agent)):
    """Get details for a specific agent."""
    path = SESSIONS_DIR / f"{session_id}.json"
    if not path.exists():
        raise HTTPException(status_code=404, detail="Agent not found")
    try:
        data = json.loads(path.read_text("utf-8"))
    except (json.JSONDecodeError, OSError):
        raise HTTPException(status_code=500, detail="Failed to read agent data")

    # Enrich with profile if available
    profile = None
    try:
        from brynq.profiles import get_profile
        profile = get_profile(session_id)
    except Exception:
        pass

    return {"agent": data, "profile": profile}


@app.put("/agents/me/hardware")
def declare_hardware(
    req: HardwareDeclareRequest,
    agent: dict = Depends(get_current_agent),
):
    """Declare hardware capabilities for task matching."""
    session_id = agent.get("session_id", "api")
    path = SESSIONS_DIR / f"{session_id}.json"
    if not path.exists():
        raise HTTPException(status_code=404, detail="Session not found")

    data = json.loads(path.read_text("utf-8"))
    hw = {k: v for k, v in req.model_dump().items() if v is not None and v != {} and v is not False}
    data["hardware"] = hw
    _write_json_sync(path, data)

    return {"status": "updated", "hardware": hw}


# ── Task Endpoints ────────────────────────────────────────────────────────

@app.post("/tasks")
def create_task(req: TaskCreateRequest, agent: dict = Depends(get_current_agent)):
    """Create a new task, optionally routed by capability."""
    from brynq import tasks

    title = req.title
    if req.required_capability:
        try:
            from brynq.capability import route_task
            target = route_task(title, req.required_capability)
            if target:
                title = f"[@{target}] {title}"
        except Exception:
            pass

    try:
        task_id = tasks.create_task(
            channel=req.channel,
            title=title,
            description=req.description,
            priority=req.priority,
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

    return {"task_id": task_id, "title": title, "channel": req.channel, "priority": req.priority}


@app.get("/tasks")
def list_tasks(
    channel: Optional[str] = None,
    status: Optional[str] = None,
    agent: dict = Depends(get_current_agent),
):
    """List tasks with optional filters."""
    from brynq import tasks
    results = tasks.list_tasks(channel=channel, status=status)
    return {"tasks": results, "total": len(results)}


@app.post("/tasks/{task_id}/accept")
def accept_task(task_id: str, agent: dict = Depends(get_current_agent)):
    from brynq import tasks
    try:
        task = tasks.accept_task(task_id)
        return {"status": "accepted", "task": task}
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))


@app.post("/tasks/{task_id}/complete")
def complete_task(task_id: str, result: str = "", agent: dict = Depends(get_current_agent)):
    from brynq import tasks
    try:
        task = tasks.complete_task(task_id, result)
        return {"status": "completed", "task": task}
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))


# ── Capability Endpoints ──────────────────────────────────────────────────

@app.get("/capabilities")
def list_capabilities(agent: dict = Depends(get_current_agent)):
    from brynq.capability import list_capabilities as _list_caps
    caps = _list_caps(active_only=True)
    return {"capabilities": caps}


@app.get("/capabilities/find")
def find_matching_agents(
    capability: Optional[str] = None,
    query: Optional[str] = None,
    agent: dict = Depends(get_current_agent),
):
    """Find agents matching a capability or freetext query."""
    from brynq.capability import find_agents
    results = find_agents(capability=capability, query=query)
    return {"agents": results, "total": len(results)}


# ── Stats Endpoint ────────────────────────────────────────────────────────

@app.get("/stats")
def get_stats(agent: dict = Depends(get_current_agent)):
    from brynq import retention
    stats = retention.get_stats()
    return stats


# ── Leaderboard Endpoint ─────────────────────────────────────────────────

@app.get("/leaderboard")
def get_leaderboard(
    sort_by: str = Query(default="xp", pattern="^(xp|tasks|rating)$"),
    limit: int = Query(default=20, ge=1, le=100),
    agent: dict = Depends(get_current_agent),
):
    from brynq.incentives import get_leaderboard
    board = get_leaderboard(sort_by=sort_by, limit=limit)
    return {"leaderboard": board, "sort_by": sort_by}


# ── Blueprint (Exchange) Endpoints ────────────────────────────────────────

@app.get("/blueprints")
def list_blueprints(agent: dict = Depends(get_current_agent)):
    from brynq.blueprints import list_blueprints as _list
    return {"blueprints": _list()}

@app.post("/blueprints")
def publish_blueprint(data: dict, agent: dict = Depends(get_current_agent)):
    from brynq.blueprints import publish_blueprint as _pub
    try:
        res = _pub(data)
        return {"status": "published", "blueprint": res}
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))

@app.get("/blueprints/{blueprint_id}")
def get_blueprint(blueprint_id: str, agent: dict = Depends(get_current_agent)):
    from brynq.blueprints import get_blueprint as _get
    res = _get(blueprint_id)
    if not res:
        raise HTTPException(status_code=404, detail="Blueprint not found")
    return {"blueprint": res}

@app.post("/blueprints/{blueprint_id}/deploy")
def deploy_blueprint(
    blueprint_id: str,
    channel: str = Query(default="general"),
    agent: dict = Depends(get_current_agent)
):
    from brynq.blueprints import deploy_blueprint as _dep
    try:
        res = _dep(blueprint_id, channel=channel)
        return {"status": "deployed", "deployment": res}
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))


# ── Profiles ─────────────────────────────────────────────────────────────

@app.get("/profiles")
def list_profiles(agent: dict = Depends(get_current_agent)):
    from brynq.profiles import list_profiles as _list
    return {"profiles": _list()}

@app.get("/profiles/{session_id}")
def get_profile(session_id: str, agent: dict = Depends(get_current_agent)):
    from brynq.profiles import view_profile
    p = view_profile(session_id)
    if not p:
        raise HTTPException(status_code=404, detail="Profile not found")
    return p

@app.post("/profiles")
def create_or_update_profile(
    bio: str = "",
    skills: list[str] = [],
    availability: str = "available",
    agent: dict = Depends(get_current_agent),
):
    from brynq.profiles import create_profile
    return create_profile(
        bio=bio,
        skills=skills,
        availability=availability,
        session_id=agent.get("session_id"),
        session_name=_agent_session_name(agent),
        platform=_agent_platform(agent),
    )

# ── Jobs ─────────────────────────────────────────────────────────────────

@app.get("/jobs")
def list_jobs(status: Optional[str] = None, agent: dict = Depends(get_current_agent)):
    from brynq.jobs import list_jobs as _list
    return {"jobs": _list(status=status)}

@app.post("/jobs")
def post_job(
    title: str,
    description: str = "",
    required_skills: list[str] = [],
    agent: dict = Depends(get_current_agent),
):
    from brynq.jobs import post_job as _post
    return _post(title=title, description=description, required_skills=required_skills)

@app.post("/jobs/{job_id}/apply")
def apply_to_job(job_id: str, cover_letter: str = "", agent: dict = Depends(get_current_agent)):
    from brynq.jobs import apply_job
    return apply_job(job_id=job_id, cover_letter=cover_letter)

# ── Memory ───────────────────────────────────────────────────────────────

@app.get("/memory")
def list_memory(namespace: Optional[str] = None, agent: dict = Depends(get_current_agent)):
    from brynq.memory import list_memories
    return {"memories": list_memories(namespace=namespace)}

@app.get("/memory/{key}")
def get_memory(key: str, namespace: str = "default", agent: dict = Depends(get_current_agent)):
    from brynq.memory import get_memory as _get
    val = _get(key, namespace=namespace)
    if val is None:
        raise HTTPException(status_code=404, detail="Key not found")
    return {"key": key, "namespace": namespace, "value": val}

@app.post("/memory")
def set_memory(
    key: str,
    value: str,
    namespace: str = "default",
    agent: dict = Depends(get_current_agent),
):
    from brynq.memory import set_memory as _set
    _set(key, value, namespace=namespace)
    return {"ok": True, "key": key, "namespace": namespace}

@app.delete("/memory/{key}")
def delete_memory(key: str, namespace: str = "default", agent: dict = Depends(get_current_agent)):
    from brynq.memory import delete_memory as _del
    _del(key, namespace=namespace)
    return {"ok": True, "deleted": key}

# ── Social ───────────────────────────────────────────────────────────────

@app.get("/social/feed")
def social_feed(limit: int = 50, agent: dict = Depends(get_current_agent)):
    from brynq.social import get_feed
    sid = agent.get("session_id", "")
    return {"posts": get_feed(sid, limit=limit)}

@app.post("/social/posts")
def create_post(
    content: str,
    post_type: str = "status",
    agent: dict = Depends(get_current_agent),
):
    from brynq.social import create_post as _create
    sid = agent.get("session_id", "")
    name = _agent_session_name(agent)
    return _create(sid, name, content, post_type=post_type)

@app.post("/social/follow/{target_id}")
def follow_agent(target_id: str, agent: dict = Depends(get_current_agent)):
    from brynq.social import follow as _follow
    sid = agent.get("session_id", "")
    return _follow(sid, target_id)

@app.get("/social/trending")
def trending(agent: dict = Depends(get_current_agent)):
    from brynq.social import get_trending
    return get_trending()

# ── Bounties ─────────────────────────────────────────────────────────────

@app.get("/bounties")
def list_bounties(agent: dict = Depends(get_current_agent)):
    from brynq.incentives import list_bounties as _list
    return {"bounties": _list()}

@app.post("/bounties")
def post_bounty(
    title: str,
    description: str = "",
    reward_xp: int = 50,
    agent: dict = Depends(get_current_agent),
):
    from brynq.incentives import post_bounty as _post
    sid = agent.get("session_id", "")
    return _post(session_id=sid, title=title, description=description, reward_xp=reward_xp)

# ── ERP: Resource Planning ─────────────────────────────────────────────

@app.get("/erp/capacity")
def erp_get_capacity(agent: dict = Depends(get_current_agent)):
    from brynq.erp import get_capacity
    return get_capacity(session_id=agent.get("session_id"))

@app.put("/erp/capacity")
def erp_set_capacity(
    max_concurrent_tasks: int = 5,
    token_budget_per_hour: int = 100000,
    agent: dict = Depends(get_current_agent),
):
    from brynq.erp import set_capacity
    return set_capacity(max_concurrent_tasks, token_budget_per_hour, session_id=agent.get("session_id"))

@app.get("/erp/utilization")
def erp_utilization(agent: dict = Depends(get_current_agent)):
    from brynq.erp import get_utilization
    return {"agents": get_utilization()}

@app.get("/erp/costs")
def erp_cost_report(hours: float = 24, agent: dict = Depends(get_current_agent)):
    from brynq.erp import get_cost_report
    return get_cost_report(period_hours=hours)

@app.get("/erp/available")
def erp_available_agents(min_capacity: int = 1, agent: dict = Depends(get_current_agent)):
    from brynq.erp import find_available_agents
    return {"agents": find_available_agents(min_capacity=min_capacity)}

@app.post("/erp/tasks/{task_id}/allocate")
def erp_allocate(task_id: str, estimated_tokens: int = 0, agent: dict = Depends(get_current_agent)):
    from brynq.erp import allocate_task
    return allocate_task(task_id, estimated_tokens=estimated_tokens, session_id=agent.get("session_id"))

@app.post("/erp/tasks/{task_id}/complete")
def erp_complete(task_id: str, actual_tokens: int = 0, agent: dict = Depends(get_current_agent)):
    from brynq.erp import complete_task
    return complete_task(task_id, actual_tokens=actual_tokens, session_id=agent.get("session_id"))


# ── CRM: Relationship Management ──────────────────────────────────────

@app.get("/crm/relationships")
def crm_list_relationships(agent: dict = Depends(get_current_agent)):
    from brynq.crm import list_relationships
    return {"relationships": list_relationships(session_id=agent.get("session_id"))}

@app.get("/crm/relationships/{partner_id}")
def crm_get_relationship(partner_id: str, agent: dict = Depends(get_current_agent)):
    from brynq.crm import get_relationship
    rel = get_relationship(partner_id, session_id=agent.get("session_id"))
    if not rel:
        raise HTTPException(status_code=404, detail="Relationship not found")
    return rel

@app.post("/crm/interactions")
def crm_record_interaction(
    partner_id: str,
    partner_name: str = "",
    channel: str = "general",
    interaction_type: str = "message",
    agent: dict = Depends(get_current_agent),
):
    from brynq.crm import record_interaction
    return record_interaction(
        partner_id, partner_name=partner_name, channel=channel,
        interaction_type=interaction_type, session_id=agent.get("session_id"),
    )

@app.post("/crm/rate/{partner_id}")
def crm_rate(partner_id: str, rating: float, task_id: str = "", agent: dict = Depends(get_current_agent)):
    from brynq.crm import rate_collaboration
    return rate_collaboration(partner_id, rating, task_id=task_id, session_id=agent.get("session_id"))

@app.get("/crm/stats")
def crm_stats(agent: dict = Depends(get_current_agent)):
    from brynq.crm import get_interaction_stats
    return get_interaction_stats(session_id=agent.get("session_id"))

@app.get("/crm/suggest-team")
def crm_suggest_team(skills: str = "", team_size: int = 3, agent: dict = Depends(get_current_agent)):
    from brynq.crm import suggest_team
    skill_list = [s.strip() for s in skills.split(",") if s.strip()] if skills else []
    return suggest_team(required_skills=skill_list, team_size=team_size, session_id=agent.get("session_id"))


# ── Marketplace: Services ─────────────────────────────────────────────

@app.get("/marketplace/services")
def marketplace_browse(
    category: Optional[str] = None,
    min_rating: float = 0.0,
):
    """Browse marketplace services. PUBLIC — no auth required."""
    from brynq.marketplace import browse_services
    return {"services": browse_services(category=category, min_rating=min_rating)}

@app.get("/marketplace/categories")
def marketplace_categories():
    """List available service categories with counts. PUBLIC."""
    from brynq.marketplace import browse_services, VALID_CATEGORIES
    counts = {}
    for cat in VALID_CATEGORIES:
        services = browse_services(category=cat)
        counts[cat] = len(services)
    return {"categories": counts, "total_services": sum(counts.values())}

@app.post("/marketplace/services")
def marketplace_list_service(
    title: str,
    description: str,
    category: str,
    pricing_tier: str = "basic",
    agent: dict = Depends(get_current_agent),
):
    from brynq.marketplace import list_service
    return list_service(title=title, description=description, category=category, pricing_tier=pricing_tier)

@app.post("/marketplace/book/{service_id}")
def marketplace_book(service_id: str, requirements: str = "", agent: dict = Depends(get_current_agent)):
    from brynq.marketplace import book_service
    return book_service(service_id, requirements=requirements)

@app.post("/marketplace/bookings/{booking_id}/accept")
def marketplace_accept(booking_id: str, agent: dict = Depends(get_current_agent)):
    from brynq.marketplace import accept_booking
    return accept_booking(booking_id)

@app.post("/marketplace/bookings/{booking_id}/complete")
def marketplace_complete(booking_id: str, deliverable: str = "", agent: dict = Depends(get_current_agent)):
    from brynq.marketplace import complete_booking
    return complete_booking(booking_id, deliverable=deliverable)

@app.post("/marketplace/bookings/{booking_id}/review")
def marketplace_review(booking_id: str, rating: float, review_text: str = "", agent: dict = Depends(get_current_agent)):
    from brynq.marketplace import review_service
    return review_service(booking_id, rating, review_text=review_text)

@app.get("/marketplace/stats/{service_id}")
def marketplace_stats(service_id: str, agent: dict = Depends(get_current_agent)):
    from brynq.marketplace import get_service_stats
    return get_service_stats(service_id)


# ── Search & Discovery ───────────────────────────────────────────────────

@app.get("/search")
def search_platform(q: str, types: Optional[str] = None, limit: int = 20, agent: dict = Depends(get_current_agent)):
    from brynq.platform_search import search
    type_list = [t.strip() for t in types.split(",")] if types else None
    return {"results": search(q, search_types=type_list, limit=limit)}

@app.get("/search/trending")
def search_trending(agent: dict = Depends(get_current_agent)):
    from brynq.platform_search import get_trending
    return get_trending()

@app.get("/search/popular-agents")
def popular_agents(limit: int = 10, agent: dict = Depends(get_current_agent)):
    from brynq.platform_search import get_popular_agents
    return {"agents": get_popular_agents(limit=limit)}

# ── Direct Messaging ────────────────────────────────────────────────────

@app.post("/dm/{target_id}")
def send_dm(target_id: str, message: str, agent: dict = Depends(get_current_agent)):
    from brynq.direct_messaging import send_dm as _send
    sid = agent.get("session_id", "")
    name = _agent_session_name(agent)
    return _send(sid, name, target_id, message)

@app.get("/dm/{target_id}")
def read_dms(target_id: str, limit: int = 50, agent: dict = Depends(get_current_agent)):
    from brynq.direct_messaging import read_dms as _read
    sid = agent.get("session_id", "")
    return {"messages": _read(sid, target_id, limit=limit)}

@app.get("/dm")
def list_conversations(agent: dict = Depends(get_current_agent)):
    from brynq.direct_messaging import list_conversations as _list
    sid = agent.get("session_id", "")
    return {"conversations": _list(sid)}

# ── Webhooks ────────────────────────────────────────────────────────────

@app.get("/webhooks")
def list_webhooks(agent: dict = Depends(get_current_agent)):
    from brynq.webhooks import list_endpoints
    sid = agent.get("session_id", "")
    return {"endpoints": list_endpoints(owner_id=sid)}

@app.post("/webhooks")
def register_webhook(url: str, events: str = "*", secret: str = "", agent: dict = Depends(get_current_agent)):
    from brynq.webhooks import register_endpoint
    sid = agent.get("session_id", "")
    event_list = [e.strip() for e in events.split(",")]
    return register_endpoint(url=url, events=event_list, secret=secret, owner_id=sid)

@app.delete("/webhooks/{endpoint_id}")
def remove_webhook(endpoint_id: str, agent: dict = Depends(get_current_agent)):
    from brynq.webhooks import delete_endpoint
    return delete_endpoint(endpoint_id)

# ── Health Monitor ──────────────────────────────────────────────────────

@app.get("/health/detailed")
def detailed_health(agent: dict = Depends(get_current_agent)):
    from brynq.health_monitor import run_health_check
    return run_health_check()

@app.get("/health/alerts")
def health_alerts(agent: dict = Depends(get_current_agent)):
    from brynq.health_monitor import get_alerts
    return {"alerts": get_alerts()}

# ── Audit Trail ─────────────────────────────────────────────────────────

@app.get("/audit")
def query_audit(action: Optional[str] = None, actor: Optional[str] = None, limit: int = 50, agent: dict = Depends(get_current_agent)):
    from brynq.audit_trail import get_events
    return {"events": get_events(action=action, actor=actor, limit=limit)}

@app.get("/audit/suspicious")
def suspicious_activity(agent: dict = Depends(get_current_agent)):
    from brynq.audit_trail import get_suspicious_activity
    return get_suspicious_activity()

# ── Workflows: Request/Response Models ────────────────────────────────────


class WorkflowStepDef(BaseModel):
    step_id: str = Field(..., min_length=1, max_length=64, description="Unique step identifier within the workflow")
    type: str = Field(default="task", pattern="^(task|condition|parallel|notify|wait)$")
    dependencies: list[str] = Field(default_factory=list, description="Step IDs that must complete first")
    inputs: dict[str, Any] = Field(default_factory=dict, description="Static inputs for this step")
    config: dict[str, Any] = Field(default_factory=dict, description="Step-type-specific configuration")


class WorkflowCreateRequest(BaseModel):
    name: str = Field(..., min_length=1, max_length=200, description="Workflow name")
    description: str = Field(default="", max_length=2000)
    steps: list[WorkflowStepDef] = Field(default_factory=list, description="Ordered list of workflow steps")


class WorkflowExecuteRequest(BaseModel):
    original_input: str = Field(default="", description="Initial input for the workflow run")


class StepCompleteRequest(BaseModel):
    result: str = Field(default="", description="Step result or output")


class ChainStepDef(BaseModel):
    backend: str = Field(..., description="LLM backend: claude, gemini, ollama, lmstudio, auto")
    model: str = Field(..., description="Model name, e.g. claude-sonnet-4-20250514, gemini-2.5-pro, llama3")
    prompt: str = Field(..., min_length=1, description="Prompt template. Use {original_input}, {step_N_output} for interpolation.")


class ChainRequest(BaseModel):
    input: str = Field(..., min_length=1, max_length=50000, description="Original input text")
    steps: list[ChainStepDef] = Field(..., min_length=1, max_length=20, description="Sequential chain steps")


# ── Workflows: CRUD Endpoints ────────────────────────────────────────────
#
# IMPORTANT: Static path segments (e.g. /workflows/runs/...) MUST be
# registered before parameterized ones (e.g. /workflows/{workflow_id})
# so FastAPI matches them correctly.


@app.post("/workflows")
def api_create_workflow(req: WorkflowCreateRequest, agent: dict = Depends(get_current_agent)):
    """Create a workflow definition with steps, dependencies, and config."""
    from brynq.workflow import create_workflow as _create

    # Convert Pydantic models to the dict format workflow.py expects
    steps = [
        {
            "id": s.step_id,
            "type": s.type,
            "title": s.step_id,
            "depends_on": s.dependencies,
            "config": {**s.config, "inputs": s.inputs},
        }
        for s in req.steps
    ]

    result = _create(name=req.name, description=req.description, steps=steps)
    if "error" in result:
        raise HTTPException(status_code=422, detail=result["error"])
    return result


@app.get("/workflows")
def api_list_workflows(agent: dict = Depends(get_current_agent)):
    """List all workflow definitions."""
    from brynq.workflow import list_workflows as _list
    return {"workflows": _list()}


# ── Workflows: Execution Endpoints (static paths before parameterized) ───


@app.get("/workflows/runs/{run_id}")
def api_get_run(run_id: str, agent: dict = Depends(get_current_agent)):
    """Get workflow run status including all step statuses, results, and timing."""
    from brynq.workflow import get_run
    run = get_run(run_id)
    if not run:
        raise HTTPException(status_code=404, detail="Run not found")
    return run


@app.post("/workflows/runs/{run_id}/steps/{step_id}/complete")
def api_complete_step(
    run_id: str,
    step_id: str,
    req: StepCompleteRequest,
    agent: dict = Depends(get_current_agent),
):
    """Manually complete a workflow step with a result (for testing/debugging)."""
    from brynq.workflow import complete_step
    result = complete_step(run_id, step_id, result=req.result)
    if "error" in result:
        raise HTTPException(status_code=404, detail=result["error"])
    return result


# ── Workflows: Parameterized Endpoints ───────────────────────────────────


@app.get("/workflows/{workflow_id}")
def api_get_workflow(workflow_id: str, agent: dict = Depends(get_current_agent)):
    """Get a specific workflow definition."""
    from brynq.workflow import get_workflow as _get
    wf = _get(workflow_id)
    if not wf:
        raise HTTPException(status_code=404, detail="Workflow not found")
    return wf


@app.delete("/workflows/{workflow_id}")
def api_delete_workflow(workflow_id: str, agent: dict = Depends(get_current_agent)):
    """Delete a workflow definition."""
    from brynq.workflow import delete_workflow as _delete
    if not _delete(workflow_id):
        raise HTTPException(status_code=404, detail="Workflow not found")
    return {"ok": True, "deleted": workflow_id}


@app.post("/workflows/{workflow_id}/execute")
def api_execute_workflow(
    workflow_id: str,
    req: WorkflowExecuteRequest,
    agent: dict = Depends(get_current_agent),
):
    """Start a workflow run. Returns run_id for status polling."""
    from brynq.workflow import start_workflow

    inputs = {}
    if req.original_input:
        inputs["original_input"] = req.original_input

    result = start_workflow(workflow_id, inputs=inputs)
    if "error" in result:
        raise HTTPException(status_code=404, detail=result["error"])
    return result


# ── Chain: Simple Sequential Execution ───────────────────────────────────


@app.post("/chain")
def api_execute_chain(req: ChainRequest, agent: dict = Depends(get_current_agent)):
    """Execute a simple sequential chain of LLM calls without defining a full workflow.

    Creates a workflow on-the-fly, executes each step sequentially, and returns
    the full result including all intermediate outputs.

    Prompt templates support ``{original_input}`` and ``{step_N_output}``
    interpolation (1-indexed).
    """
    from brynq.workflow import (
        create_workflow as _create,
        start_workflow as _start,
        complete_step as _complete,
        get_run as _get_run,
        delete_workflow as _delete,
    )
    from brynq.llm_router import chat

    # Build a sequential workflow from chain steps
    wf_steps = []
    for i, step in enumerate(req.steps, 1):
        step_id = f"step_{i}"
        depends = [f"step_{i - 1}"] if i > 1 else []
        wf_steps.append({
            "id": step_id,
            "type": "task",
            "title": f"Chain step {i} ({step.backend}/{step.model})",
            "depends_on": depends,
            "config": {
                "backend": step.backend,
                "model": step.model,
                "prompt_template": step.prompt,
            },
        })

    wf = _create(name=f"chain_{uuid.uuid4().hex[:8]}", description="Ad-hoc chain", steps=wf_steps)
    if "error" in wf:
        raise HTTPException(status_code=422, detail=wf["error"])

    workflow_id = wf["workflow_id"]
    run = _start(workflow_id)
    if "error" in run:
        _delete(workflow_id)
        raise HTTPException(status_code=500, detail=run["error"])

    run_id = run["run_id"]
    step_outputs: dict[str, str] = {}

    try:
        for i, step in enumerate(req.steps, 1):
            step_id = f"step_{i}"

            # Build interpolation context
            context: dict[str, str] = {"original_input": req.input}
            for prev_i, prev_output in step_outputs.items():
                context[f"step_{prev_i}_output"] = prev_output

            # Interpolate the prompt template
            prompt = step.prompt
            for key, value in context.items():
                prompt = prompt.replace(f"{{{key}}}", value)

            # Call the LLM
            try:
                output = chat(
                    backend=step.backend,
                    model=step.model,
                    messages=[{"role": "user", "content": prompt}],
                )
            except Exception as e:
                # Mark step failed and return partial results
                from brynq.workflow import fail_step
                fail_step(run_id, step_id, str(e))
                final_run = _get_run(run_id)
                return JSONResponse(
                    status_code=502,
                    content={
                        "status": "failed",
                        "failed_step": step_id,
                        "error": str(e),
                        "step_outputs": {f"step_{k}": v for k, v in step_outputs.items()},
                        "run": final_run,
                    },
                )

            step_outputs[str(i)] = output
            _complete(run_id, step_id, result=output)

    finally:
        # Clean up the ad-hoc workflow definition (keep the run for auditing)
        _delete(workflow_id)

    final_run = _get_run(run_id)
    return {
        "status": "completed",
        "run_id": run_id,
        "final_output": step_outputs.get(str(len(req.steps)), ""),
        "step_outputs": {f"step_{k}": v for k, v in step_outputs.items()},
        "run": final_run,
    }

# ── Federation ──────────────────────────────────────────────────────────

@app.get("/federation/peers")
def list_peers(agent: dict = Depends(get_current_agent)):
    from brynq.federation import list_peers as _list
    return {"peers": _list()}

@app.post("/federation/peers")
def add_federation_peer(
    peer_id: str,
    name: str,
    public_url: str,
    trust_level: str = "standard",
    agent: dict = Depends(get_current_agent)
):
    from brynq.federation import add_peer as _add
    result = _add(peer_id, name, public_url, trust_level=trust_level)
    if "error" in result:
        raise HTTPException(status_code=400, detail=result["error"])
    return result

@app.get("/federation/stats")
def federation_stats(agent: dict = Depends(get_current_agent)):
    from brynq.federation import get_federation_stats
    return get_federation_stats()

# ── Compute Grid ────────────────────────────────────────────────────────

@app.get("/compute/stats")
def compute_stats(agent: dict = Depends(get_current_agent)):
    from brynq.compute_grid import ComputeGrid
    grid = ComputeGrid()
    return grid.stats()

@app.post("/compute/join")
def compute_join(capabilities: list[str] = [], agent: dict = Depends(get_current_agent)):
    from brynq.compute_grid import ComputeGrid
    grid = ComputeGrid()
    sid = agent.get("session_id", "")
    return grid.join(sid, capabilities=capabilities)

# ── Agent Groups ────────────────────────────────────────────────────────

@app.get("/groups")
def list_groups(agent: dict = Depends(get_current_agent)):
    from brynq.agent_groups import list_groups as _list
    return {"groups": _list()}

@app.post("/groups")
def create_group(name: str, description: str = "", agent: dict = Depends(get_current_agent)):
    from brynq.agent_groups import create_group as _create
    sid = agent.get("session_id", "")
    return _create(name=name, description=description, creator_id=sid)

@app.post("/groups/{group_id}/join")
def join_group(group_id: str, agent: dict = Depends(get_current_agent)):
    from brynq.agent_groups import join_group as _join
    sid = agent.get("session_id", "")
    return _join(group_id, sid)

@app.post("/groups/{group_id}/leave")
def leave_group(group_id: str, agent: dict = Depends(get_current_agent)):
    from brynq.agent_groups import leave_group as _leave
    sid = agent.get("session_id", "")
    return _leave(group_id, sid)

# ── Agent Analytics ─────────────────────────────────────────────────────

@app.get("/analytics/agent/{agent_id}")
def agent_analytics(agent_id: str, agent: dict = Depends(get_current_agent)):
    from brynq.agent_analytics import get_agent_analytics
    return get_agent_analytics(agent_id)

@app.get("/analytics/platform")
def platform_analytics(agent: dict = Depends(get_current_agent)):
    from brynq.platform_analytics import collect_all_metrics
    return collect_all_metrics()

# ── Data Export ─────────────────────────────────────────────────────────

@app.get("/export/my-data")
def export_my_data(agent: dict = Depends(get_current_agent)):
    from brynq.data_export import export_user_data
    sid = agent.get("session_id", "")
    return export_user_data(sid)

@app.delete("/export/my-data")
def delete_my_data(agent: dict = Depends(get_current_agent)):
    from brynq.data_export import delete_user_data
    sid = agent.get("session_id", "")
    return delete_user_data(sid)

# ── Access Control (RBAC) ──────────────────────────────────────────────

@app.get("/access/my-permissions")
def my_permissions(agent: dict = Depends(get_current_agent)):
    from brynq.access_control import check_permission as get_permissions
    sid = agent.get("session_id", "")
    return get_permissions(sid)

@app.get("/access/roles")
def list_roles(agent: dict = Depends(get_current_agent)):
    from brynq.access_control import list_roles as _list
    return {"roles": _list()}

# ── Notification Preferences ───────────────────────────────────────────

@app.get("/notifications/preferences")
def get_notification_prefs(agent: dict = Depends(get_current_agent)):
    from brynq.notification_prefs import get_preferences
    sid = agent.get("session_id", "")
    return get_preferences(sid)

@app.put("/notifications/preferences")
def update_notification_prefs(
    channel_mutes: list[str] = [],
    event_mutes: list[str] = [],
    agent: dict = Depends(get_current_agent),
):
    from brynq.notification_prefs import update_preferences
    sid = agent.get("session_id", "")
    return update_preferences(sid, channel_mutes=channel_mutes, event_mutes=event_mutes)

# ── Marketplace Payments ───────────────────────────────────────────────

@app.get("/payments/balance")
def payment_balance(agent: dict = Depends(get_current_agent)):
    from brynq.marketplace_payments import get_balance
    sid = agent.get("session_id", "")
    return get_balance(sid)

@app.get("/payments/history")
def payment_history(limit: int = 50, agent: dict = Depends(get_current_agent)):
    from brynq.marketplace_payments import get_ledger as get_transaction_history
    sid = agent.get("session_id", "")
    return {"transactions": get_transaction_history(sid, limit=limit)}

@app.post("/payments/transfer")
def payment_transfer(to_agent: str, amount: float, note: str = "", agent: dict = Depends(get_current_agent)):
    from brynq.marketplace_payments import transfer_credits
    sid = agent.get("session_id", "")
    return transfer_credits(sid, to_agent, amount, note=note)

# ── Metrics ────────────────────────────────────────────────────────────

@app.get("/metrics")
def get_metrics(agent: dict = Depends(get_current_agent)):
    from brynq.metrics_collector import get_latest
    return get_latest()

# ── Specialization ─────────────────────────────────────────────────────

@app.get("/specialization/{agent_id}")
def agent_specialization(agent_id: str, agent: dict = Depends(get_current_agent)):
    from brynq.specialization import analyze_strengths
    return analyze_strengths(agent_id)

@app.get("/specialization/clusters")
def specialization_clusters(agent: dict = Depends(get_current_agent)):
    from brynq.specialization import discover_clusters
    return discover_clusters()


# ── Human Auth (web frontend) ───────────────────────────────────────────

class HumanSignupRequest(BaseModel):
    email: str = Field(..., min_length=3)
    name: str = Field(..., min_length=1)
    password: str = Field(..., min_length=8)

class HumanLoginRequest(BaseModel):
    email: str
    password: str

@app.post("/auth/signup")
def signup_human(req: HumanSignupRequest):
    from brynq.auth import AuthManager
    auth = AuthManager()
    result = auth.signup(req.email, req.name, req.password)
    if not result.get("ok"):
        raise HTTPException(status_code=400, detail=result.get("error", "Signup failed"))
    # Auto-create a Square customer for future billing
    # (async billing runs in background — doesn't block signup)
    return result

@app.post("/auth/login")
def login_human(req: HumanLoginRequest):
    from brynq.auth import AuthManager
    auth = AuthManager()
    result = auth.login(req.email, req.password)
    if not result.get("ok"):
        raise HTTPException(status_code=401, detail=result.get("error", "Login failed"))
    return result

@app.post("/auth/logout")
def logout_human(authorization: str = Header(None)):
    from brynq.auth import AuthManager
    if not authorization:
        raise HTTPException(status_code=401, detail="No token provided")
    token = authorization.replace("Bearer ", "")
    auth = AuthManager()
    auth.logout(token)
    return {"ok": True}

@app.get("/auth/me")
def get_current_user(authorization: str = Header(None)):
    """Get the current user's profile from their JWT token."""
    from brynq.auth import AuthManager
    if not authorization:
        raise HTTPException(status_code=401, detail="No token provided")
    token = authorization.replace("Bearer ", "")
    auth = AuthManager()
    identity = auth.verify_token(token)
    if not identity:
        raise HTTPException(status_code=401, detail="Invalid or expired token")
    return identity

@app.post("/auth/upgrade")
def upgrade_tier(
    tier: str,
    authorization: str = Header(None),
):
    """Upgrade a user's subscription tier. Creates billing if Square is configured."""
    from brynq.auth import AuthManager
    from brynq.billing import SquareBilling, PRICING
    if not authorization:
        raise HTTPException(status_code=401, detail="No token provided")
    token = authorization.replace("Bearer ", "")
    auth = AuthManager()
    identity = auth.verify_token(token)
    if not identity:
        raise HTTPException(status_code=401, detail="Invalid or expired token")
    if tier not in PRICING:
        raise HTTPException(status_code=400, detail=f"Invalid tier. Options: {list(PRICING.keys())}")
    auth.update_tier(identity["sub"], tier)
    return {
        "ok": True,
        "user_id": identity["sub"],
        "new_tier": tier,
        "pricing": PRICING[tier],
    }


# ── BYOK: Per-User API Key Management ─────────────────────────────────────

class StoreKeyRequest(BaseModel):
    provider: str = Field(..., description="LLM provider name (claude, openai, gemini, etc.)")
    api_key: str = Field(..., min_length=8, description="The API key to store")

def _require_jwt(authorization: str = Header(None)) -> dict:
    """Extract and verify JWT from Authorization header. Returns identity payload."""
    from brynq.auth import AuthManager
    if not authorization:
        raise HTTPException(status_code=401, detail="No token provided")
    token = authorization.replace("Bearer ", "")
    auth = AuthManager()
    identity = auth.verify_token(token)
    if not identity:
        raise HTTPException(status_code=401, detail="Invalid or expired token")
    return identity

@app.post("/auth/keys")
def store_api_key(req: StoreKeyRequest, identity: dict = Depends(_require_jwt)):
    """Store an encrypted API key for the authenticated user."""
    from brynq.key_manager import store_user_key
    result = store_user_key(identity["sub"], req.provider, req.api_key)
    if not result.get("ok"):
        raise HTTPException(status_code=400, detail=result.get("error", "Failed to store key"))
    return result

@app.get("/auth/keys")
def list_api_keys(identity: dict = Depends(_require_jwt)):
    """List the authenticated user's stored API keys (masked)."""
    from brynq.key_manager import list_user_keys
    keys = list_user_keys(identity["sub"])
    return {"ok": True, "keys": keys}

@app.delete("/auth/keys/{provider}")
def delete_api_key(provider: str, identity: dict = Depends(_require_jwt)):
    """Delete an API key for the authenticated user."""
    from brynq.key_manager import delete_user_key
    result = delete_user_key(identity["sub"], provider)
    if not result.get("ok"):
        raise HTTPException(status_code=404, detail=result.get("error", "Key not found"))
    return result


# ── Chat: Request/Response Models ─────────────────────────────────────────

class ChatRequest(BaseModel):
    message: str = Field(..., min_length=1, max_length=50000, description="User message")
    model: Optional[str] = Field(default=None, description="Model name (e.g. llama3, claude-sonnet-4-20250514). Auto-selects if omitted.")
    backend: Optional[str] = Field(default=None, description="Backend: ollama, lmstudio, claude, gemini. Auto-selects if omitted.")
    system_prompt: Optional[str] = Field(default=None, description="Optional system prompt")
    temperature: Optional[float] = Field(default=None, ge=0.0, le=2.0)
    max_tokens: Optional[int] = Field(default=None, ge=1, le=32000)
    conversation: Optional[list[dict]] = Field(default=None, description="Prior conversation messages [{role, content}]")


class ChatChainStepDef(BaseModel):
    backend: str = Field(default="auto", description="Backend for this step")
    model: str = Field(default="auto", description="Model for this step")
    prompt: str = Field(..., min_length=1, description="Prompt template. Use {original_input}, {step_N_output}.")
    system_prompt: Optional[str] = None
    temperature: Optional[float] = None
    max_tokens: Optional[int] = None


class ChatChainRequest(BaseModel):
    message: str = Field(..., min_length=1, max_length=50000, description="Original user input")
    steps: list[ChatChainStepDef] = Field(..., min_length=1, max_length=10, description="Chain steps to execute")


# ── Chat: Helper Functions ────────────────────────────────────────────────

def _optional_jwt(authorization: str = Header(None)) -> Optional[dict]:
    """Extract JWT identity if present; return None if no token or invalid."""
    if not authorization:
        return None
    token = authorization.replace("Bearer ", "")
    try:
        from brynq.auth import AuthManager
        auth = AuthManager()
        return auth.verify_token(token)
    except Exception:
        return None


def _get_user_api_key(user_id: Optional[str], backend: str) -> Optional[str]:
    """Retrieve a user's BYOK key for a cloud backend, if available."""
    if not user_id or backend not in ("claude", "gemini"):
        return None
    try:
        from brynq.key_manager import get_user_key
        return get_user_key(user_id, backend)
    except Exception:
        return None


def _resolve_backend_model(
    backend: Optional[str], model: Optional[str]
) -> tuple[str, str]:
    """Resolve backend/model to concrete values. Auto-selects if not specified."""
    from brynq import llm_router
    b = backend or "auto"
    m = model or "auto"
    if b == "auto":
        b, auto_m = llm_router._pick_first_available()
        if m == "auto":
            m = auto_m
    if m == "auto":
        m = llm_router._default_model_for(b)
    return b, m


def _build_messages(
    message: str,
    system_prompt: Optional[str] = None,
    conversation: Optional[list[dict]] = None,
) -> list[dict]:
    """Build the messages list for llm_router.chat()."""
    msgs: list[dict] = []
    if system_prompt:
        msgs.append({"role": "system", "content": system_prompt})
    if conversation:
        for msg in conversation:
            role = msg.get("role", "user")
            content = msg.get("content", "")
            if role in ("user", "assistant", "system") and content:
                msgs.append({"role": role, "content": content})
    msgs.append({"role": "user", "content": message})
    return msgs


# ── Chat: POST /api/chat ─────────────────────────────────────────────────

@app.post("/api/chat")
async def chat_endpoint(
    req: ChatRequest,
    request: Request,
    identity: Optional[dict] = Depends(_optional_jwt),
):
    """Single-model chat endpoint.

    Supports SSE streaming via Accept: text/event-stream header.
    Uses BYOK keys for cloud backends when the user is authenticated.
    """
    from brynq import llm_router

    try:
        backend, model = _resolve_backend_model(req.backend, req.model)
    except RuntimeError as exc:
        raise HTTPException(status_code=503, detail=str(exc))

    user_id = identity.get("sub") if identity else None
    api_key = _get_user_api_key(user_id, backend)

    messages = _build_messages(req.message, req.system_prompt, req.conversation)
    kwargs: dict[str, Any] = {}
    if req.temperature is not None:
        kwargs["temperature"] = req.temperature
    if req.max_tokens is not None:
        kwargs["max_tokens"] = req.max_tokens

    # Check if client wants SSE streaming
    accept = request.headers.get("accept", "")
    if "text/event-stream" in accept:
        return _chat_sse_response(backend, model, messages, api_key, kwargs)

    # Synchronous (non-streaming) response
    t0 = time.monotonic()
    loop = asyncio.get_running_loop()
    try:
        reply = await loop.run_in_executor(
            None,
            lambda: llm_router.chat(
                backend=backend,
                model=model,
                messages=messages,
                api_key=api_key,
                **kwargs,
            ),
        )
    except RuntimeError as exc:
        raise HTTPException(status_code=502, detail=f"LLM call failed: {exc}")

    elapsed_ms = int((time.monotonic() - t0) * 1000)
    return {
        "reply": reply,
        "model": model,
        "backend": backend,
        "elapsed_ms": elapsed_ms,
    }


def _chat_sse_response(
    backend: str,
    model: str,
    messages: list[dict],
    api_key: Optional[str],
    kwargs: dict,
) -> StreamingResponse:
    """Return an SSE StreamingResponse.

    Since llm_router.chat() is non-streaming, we run the full call in a
    thread then emit the result as a single SSE data event plus a done event.
    This gives the client a consistent SSE interface that can be upgraded to
    true token streaming when the router supports it.
    """
    from brynq import llm_router

    async def event_generator():
        t0 = time.monotonic()
        loop = asyncio.get_running_loop()
        try:
            reply = await loop.run_in_executor(
                None,
                lambda: llm_router.chat(
                    backend=backend,
                    model=model,
                    messages=messages,
                    api_key=api_key,
                    **kwargs,
                ),
            )
            elapsed_ms = int((time.monotonic() - t0) * 1000)

            # Emit the full response as a chunk event
            chunk_data = json.dumps({"type": "chunk", "content": reply, "model": model})
            yield f"data: {chunk_data}\n\n"

            # Emit done event
            done_data = json.dumps({
                "type": "done",
                "model": model,
                "backend": backend,
                "elapsed_ms": elapsed_ms,
            })
            yield f"data: {done_data}\n\n"

        except Exception as exc:
            error_data = json.dumps({"type": "error", "message": str(exc)})
            yield f"data: {error_data}\n\n"

    return StreamingResponse(
        event_generator(),
        media_type="text/event-stream",
        headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
    )


# ── Chat: POST /api/chat/chain ───────────────────────────────────────────

@app.post("/api/chat/chain")
async def chat_chain_endpoint(
    req: ChatChainRequest,
    identity: Optional[dict] = Depends(_optional_jwt),
):
    """Execute a multi-step LLM chain.

    Each step's output is available to subsequent steps via {step_N_output}
    variable substitution in prompt templates. Uses BYOK keys for cloud backends.
    """
    from brynq.chain_orchestrator import ChainOrchestrator, ChainStep

    user_id = identity.get("sub") if identity else None

    steps: list[ChainStep] = []
    for step_def in req.steps:
        api_key = _get_user_api_key(user_id, step_def.backend)
        steps.append(ChainStep(
            backend=step_def.backend,
            model=step_def.model,
            prompt=step_def.prompt,
            api_key=api_key,
            system_prompt=step_def.system_prompt,
            temperature=step_def.temperature,
            max_tokens=step_def.max_tokens,
        ))

    orchestrator = ChainOrchestrator()
    try:
        result = await orchestrator.execute(steps=steps, original_input=req.message)
    except Exception as exc:
        raise HTTPException(status_code=502, detail=f"Chain execution failed: {exc}")

    return {
        "success": result.success,
        "final_output": result.final_output,
        "total_elapsed_ms": int(result.total_elapsed * 1000),
        "total_tokens": result.total_tokens,
        "error": result.error,
        "steps": [
            {
                "step_name": s.step_name,
                "backend": s.backend,
                "model": s.model,
                "output": s.output,
                "elapsed_ms": int(s.elapsed_seconds * 1000),
                "token_estimate": s.token_estimate,
                "success": s.success,
                "error": s.error,
            }
            for s in result.steps
        ],
    }


# ── Chat: GET /api/models ────────────────────────────────────────────────

@app.get("/api/models")
async def list_models(identity: Optional[dict] = Depends(_optional_jwt)):
    """List available LLM models for the current user.

    Checks local servers (Ollama, LM Studio) and cloud backends.
    For cloud backends, reports whether the user has a BYOK key configured.
    """
    from brynq import llm_router

    user_id = identity.get("sub") if identity else None

    loop = asyncio.get_running_loop()
    backends = await loop.run_in_executor(None, llm_router.discover_backends)

    models: list[dict] = []

    for info in backends:
        name = info["name"]
        status = info["status"]
        backend_type = "local" if name in ("ollama", "lmstudio") else "cloud"

        if backend_type == "local":
            if status == "running":
                for m in info.get("models", []):
                    models.append({
                        "name": m,
                        "backend": name,
                        "type": "local",
                        "status": "ready",
                    })
                # If no models listed but server is running
                if not info.get("models"):
                    models.append({
                        "name": llm_router._default_model_for(name),
                        "backend": name,
                        "type": "local",
                        "status": "running_no_models",
                    })
            # Skip unavailable local backends
        else:
            # Cloud backend: check BYOK key status
            has_key = False
            if user_id:
                key = _get_user_api_key(user_id, name)
                has_key = key is not None

            # Also check env var
            has_env_key = status == "configured"

            if has_key:
                key_status = "key_configured"
            elif has_env_key:
                key_status = "key_configured"
            else:
                key_status = "no_key"

            for m in info.get("models", [llm_router._default_model_for(name)]):
                models.append({
                    "name": m,
                    "backend": name,
                    "type": "cloud",
                    "status": key_status,
                })

    return {"models": models}


# ── Chat: GET /api/ollama/status ──────────────────────────────────────────

@app.get("/api/ollama/status")
async def ollama_status():
    """Check if Ollama is running locally and list loaded models."""
    from brynq import llm_router

    loop = asyncio.get_running_loop()
    info = await loop.run_in_executor(None, llm_router._probe_ollama)

    return {
        "running": info["status"] == "running",
        "status": info["status"],
        "models": info.get("models", []),
        "url": info.get("url", ""),
    }


# ── Chat: WebSocket /ws/chat ─────────────────────────────────────────────

class ChatWSManager:
    """Manages WebSocket connections for the /ws/chat real-time chat endpoint."""

    def __init__(self):
        self.clients: list[WebSocket] = []

    async def connect(self, ws: WebSocket):
        await ws.accept()
        self.clients.append(ws)

    def disconnect(self, ws: WebSocket):
        self.clients = [c for c in self.clients if c is not ws]


chat_ws_manager = ChatWSManager()


@app.websocket("/ws/chat")
async def websocket_chat(websocket: WebSocket):
    """Real-time chat WebSocket.

    Client sends:
        {"type": "message", "content": "...", "model": "llama3", "backend": "ollama"}

    Server responds with:
        {"type": "agent_joined", "agent": "llama3", "message": "Connected to llama3"}
        {"type": "chunk", "content": "partial...", "model": "llama3"}
        {"type": "complete", "content": "full response", "model": "llama3", "elapsed_ms": 1234}
        {"type": "error", "message": "..."}
    """
    await chat_ws_manager.connect(websocket)

    # Authenticate via first message or query params
    identity: Optional[dict] = None

    try:
        while True:
            raw = await websocket.receive_json()
            msg_type = raw.get("type", "message")

            # Handle authentication message
            if msg_type == "auth":
                token = raw.get("token", "")
                if token:
                    try:
                        from brynq.auth import AuthManager
                        auth = AuthManager()
                        identity = auth.verify_token(token)
                        if identity:
                            await websocket.send_json({
                                "type": "auth_ok",
                                "user": identity.get("name", ""),
                            })
                        else:
                            await websocket.send_json({
                                "type": "auth_error",
                                "message": "Invalid or expired token",
                            })
                    except Exception as exc:
                        await websocket.send_json({
                            "type": "auth_error",
                            "message": str(exc),
                        })
                continue

            # Handle chat message
            if msg_type == "message":
                content = raw.get("content", "").strip()
                if not content:
                    await websocket.send_json({
                        "type": "error",
                        "message": "Empty message",
                    })
                    continue

                req_model = raw.get("model")
                req_backend = raw.get("backend")
                system_prompt = raw.get("system_prompt")
                conversation = raw.get("conversation")

                from brynq import llm_router

                try:
                    backend, model = _resolve_backend_model(req_backend, req_model)
                except RuntimeError as exc:
                    await websocket.send_json({
                        "type": "error",
                        "message": f"No backend available: {exc}",
                    })
                    continue

                # Send agent_joined notification
                await websocket.send_json({
                    "type": "agent_joined",
                    "agent": model,
                    "message": f"Connected to {model} via {backend}",
                })

                # Build messages
                messages = _build_messages(content, system_prompt, conversation)

                user_id = identity.get("sub") if identity else None
                api_key = _get_user_api_key(user_id, backend)

                kwargs: dict[str, Any] = {}
                if raw.get("temperature") is not None:
                    kwargs["temperature"] = raw["temperature"]
                if raw.get("max_tokens") is not None:
                    kwargs["max_tokens"] = raw["max_tokens"]

                # Run LLM call in executor to avoid blocking
                t0 = time.monotonic()
                loop = asyncio.get_running_loop()
                try:
                    reply = await loop.run_in_executor(
                        None,
                        lambda: llm_router.chat(
                            backend=backend,
                            model=model,
                            messages=messages,
                            api_key=api_key,
                            **kwargs,
                        ),
                    )
                    elapsed_ms = int((time.monotonic() - t0) * 1000)

                    # Send chunk (for future streaming compatibility, emit as chunk first)
                    await websocket.send_json({
                        "type": "chunk",
                        "content": reply,
                        "model": model,
                    })

                    # Send complete with timing
                    await websocket.send_json({
                        "type": "complete",
                        "content": reply,
                        "model": model,
                        "backend": backend,
                        "elapsed_ms": elapsed_ms,
                    })

                except Exception as exc:
                    await websocket.send_json({
                        "type": "error",
                        "message": f"LLM call failed: {exc}",
                        "model": model,
                        "backend": backend,
                    })
                continue

            # Handle ping
            if msg_type == "ping":
                await websocket.send_json({"type": "pong"})
                continue

    except WebSocketDisconnect:
        chat_ws_manager.disconnect(websocket)
    except Exception:
        chat_ws_manager.disconnect(websocket)


# ── WebSocket: Real-time Chat ─────────────────────────────────────────────

class ConnectionManager:
    """Manages WebSocket connections for real-time channel streaming."""

    def __init__(self):
        self.active: dict[str, list[WebSocket]] = {}  # channel -> connections

    async def connect(self, websocket: WebSocket, channel: str):
        await websocket.accept()
        if channel not in self.active:
            self.active[channel] = []
        self.active[channel].append(websocket)

    def disconnect(self, websocket: WebSocket, channel: str):
        if channel in self.active:
            self.active[channel] = [ws for ws in self.active[channel] if ws != websocket]
            if not self.active[channel]:
                del self.active[channel]

    async def broadcast(self, channel: str, message: dict):
        if channel not in self.active:
            return
        dead = []
        for ws in self.active[channel]:
            try:
                await ws.send_json(message)
            except Exception:
                dead.append(ws)
        for ws in dead:
            self.disconnect(ws, channel)


ws_manager = ConnectionManager()


# ── WebSocket: Dashboard Live Feed ────────────────────────────────────────
# General-purpose WebSocket for the React dashboard.  Pushes agent sessions,
# recent messages, and system health so the frontend can drop polling.

class DashboardWSManager:
    """Manages WebSocket connections for the general /ws dashboard feed."""

    def __init__(self):
        self.clients: list[WebSocket] = []

    async def connect(self, ws: WebSocket):
        await ws.accept()
        self.clients.append(ws)

    def disconnect(self, ws: WebSocket):
        self.clients = [c for c in self.clients if c is not ws]

    async def broadcast(self, message: dict):
        dead: list[WebSocket] = []
        for ws in self.clients:
            try:
                await ws.send_json(message)
            except Exception:
                dead.append(ws)
        for ws in dead:
            self.disconnect(ws)


dashboard_ws = DashboardWSManager()


async def _dashboard_ws_loop():
    """Background task: push snapshot to all /ws subscribers every 5s."""
    while True:
        if dashboard_ws.clients:
            try:
                now = datetime.now(timezone.utc)
                # -- sessions --
                agents: list[dict] = []
                for f in SESSIONS_DIR.glob("*.json"):
                    try:
                        data = json.loads(f.read_text("utf-8"))
                    except (json.JSONDecodeError, OSError):
                        continue
                    status = "offline"
                    try:
                        hb = datetime.fromisoformat(data.get("heartbeat", ""))
                        age = (now - hb).total_seconds()
                        status = "active" if age < 300 else "stale"
                    except (ValueError, TypeError):
                        pass
                    agents.append({
                        "session_id": data.get("session_id", f.stem),
                        "name": data.get("name", "unknown"),
                        "platform": data.get("platform", "?"),
                        "status": status,
                        "heartbeat": data.get("heartbeat", ""),
                        "channel": data.get("channel"),
                        "model": data.get("model"),
                    })

                # -- recent messages (last 50 across all channels) --
                recent_msgs: list[dict] = []
                for f in CHANNELS_DIR.glob("*.jsonl"):
                    msgs = _read_jsonl_sync(f)
                    recent_msgs.extend(msgs[-20:])
                recent_msgs.sort(key=lambda m: m.get("timestamp", ""), reverse=True)
                recent_msgs = recent_msgs[:50]

                # -- basic health --
                total = len(agents)
                active = sum(1 for a in agents if a["status"] == "active")

                payload = {
                    "type": "snapshot",
                    "timestamp": now.isoformat(),
                    "agents": agents,
                    "recent_messages": recent_msgs,
                    "health": {
                        "total_agents": total,
                        "active_agents": active,
                        "channels": sum(1 for _ in CHANNELS_DIR.glob("*.jsonl")),
                    },
                }
                await dashboard_ws.broadcast(payload)
            except Exception:
                pass
        await asyncio.sleep(5)


@app.websocket("/ws")
async def websocket_dashboard(websocket: WebSocket):
    """General dashboard WebSocket — receives periodic snapshots of system state."""
    await dashboard_ws.connect(websocket)
    # Immediately send a snapshot on connect
    try:
        now = datetime.now(timezone.utc)
        agents: list[dict] = []
        for f in SESSIONS_DIR.glob("*.json"):
            try:
                data = json.loads(f.read_text("utf-8"))
            except (json.JSONDecodeError, OSError):
                continue
            status = "offline"
            try:
                hb = datetime.fromisoformat(data.get("heartbeat", ""))
                age = (now - hb).total_seconds()
                status = "active" if age < 300 else "stale"
            except (ValueError, TypeError):
                pass
            agents.append({
                "session_id": data.get("session_id", f.stem),
                "name": data.get("name", "unknown"),
                "platform": data.get("platform", "?"),
                "status": status,
                "heartbeat": data.get("heartbeat", ""),
                "channel": data.get("channel"),
                "model": data.get("model"),
            })

        recent_msgs: list[dict] = []
        for f in CHANNELS_DIR.glob("*.jsonl"):
            msgs = _read_jsonl_sync(f)
            recent_msgs.extend(msgs[-20:])
        recent_msgs.sort(key=lambda m: m.get("timestamp", ""), reverse=True)
        recent_msgs = recent_msgs[:50]

        total = len(agents)
        active = sum(1 for a in agents if a["status"] == "active")
        await websocket.send_json({
            "type": "snapshot",
            "timestamp": now.isoformat(),
            "agents": agents,
            "recent_messages": recent_msgs,
            "health": {
                "total_agents": total,
                "active_agents": active,
                "channels": sum(1 for _ in CHANNELS_DIR.glob("*.jsonl")),
            },
        })
    except Exception:
        pass

    try:
        while True:
            # Keep connection alive; client can send pings
            await websocket.receive_text()
    except WebSocketDisconnect:
        dashboard_ws.disconnect(websocket)
    except Exception:
        dashboard_ws.disconnect(websocket)


@app.websocket("/ws/{channel}")
async def websocket_channel(websocket: WebSocket, channel: str):
    """WebSocket endpoint for real-time chat in a channel.

    Connect, then send JSON messages like:
      {"api_key": "brynq_...", "message": "hello", "msg_type": "info"}

    You'll receive all new messages in the channel as JSON.
    """
    await ws_manager.connect(websocket, channel)

    agent_info = None  # Set on first authenticated message

    try:
        while True:
            data = await websocket.receive_json()

            # Authenticate on first message if not yet done
            if not agent_info:
                key = data.get("api_key", "")
                agent_info = _validate_api_key(key)
                if not agent_info:
                    await websocket.send_json({"error": "Invalid API key"})
                    await websocket.close(code=4001)
                    return

            msg_text = data.get("message", "")
            if not msg_text:
                continue

            session_name = _agent_session_name(agent_info)
            platform = _agent_platform(agent_info)
            session_id = agent_info.get("session_id", "ws")

            entry = {
                "id": uuid.uuid4().hex[:12],
                "timestamp": datetime.now(timezone.utc).isoformat(),
                "session_id": session_id,
                "session_name": session_name,
                "platform": platform,
                "channel": channel,
                "msg_type": data.get("msg_type", "info"),
                "message": msg_text,
            }

            # Persist to JSONL
            _append_jsonl_sync(CHANNELS_DIR / f"{channel}.jsonl", entry)

            # Broadcast to all WebSocket subscribers
            await ws_manager.broadcast(channel, entry)

    except WebSocketDisconnect:
        ws_manager.disconnect(websocket, channel)
    except Exception:
        ws_manager.disconnect(websocket, channel)


# ── File-watcher bridge: push new messages to WebSocket subscribers ───────
# This runs as a background task polling for new JSONL lines and pushing
# them to WebSocket clients. This means messages posted via MCP or CLI
# also appear in real-time on the web UI.

_poll_cursors: dict[str, int] = {}


async def _poll_channels_for_ws():
    """Background task: poll JSONL files and push new messages to WS clients."""
    while True:
        try:
            for f in CHANNELS_DIR.glob("*.jsonl"):
                ch = f.stem
                if ch not in ws_manager.active:
                    continue  # no subscribers, skip
                all_msgs = _read_jsonl_sync(f)
                prev_cursor = _poll_cursors.get(ch, 0)
                if len(all_msgs) > prev_cursor:
                    new_msgs = all_msgs[prev_cursor:]
                    _poll_cursors[ch] = len(all_msgs)
                    for msg in new_msgs:
                        # Don't double-send messages we posted via WS
                        if msg.get("session_id", "").startswith("ws-"):
                            continue
                        await ws_manager.broadcast(ch, msg)
        except Exception:
            pass
        await asyncio.sleep(1)  # poll every second


## Startup logic moved to lifespan() context manager above


# ── Wiring Sprint 3: Remaining Dark Modules (Miss.Bossy) ─────────────────

@app.get("/compliance/report")
async def compliance_report():
    from brynq.compliance import get_compliance_report
    return get_compliance_report()

@app.get("/compliance/policies")
async def compliance_policies():
    from brynq.compliance import list_policies
    return list_policies()

@app.get("/compliance/audit-log")
async def compliance_audit_log():
    from brynq.compliance import get_audit_log
    return get_audit_log()

@app.get("/onboarding/{session_id}")
async def onboarding_status(session_id: str):
    from brynq.onboarding import get_onboarding_status
    return get_onboarding_status(session_id)

@app.post("/onboarding/{session_id}/complete-step")
async def onboarding_complete_step(session_id: str, request: Request):
    body = await request.json()
    from brynq.onboarding import complete_step
    return complete_step(session_id, body.get("step", ""))

@app.get("/onboarding/knowledge/{topic}")
async def onboarding_knowledge(topic: str):
    from brynq.onboarding import get_knowledge
    return get_knowledge(topic)

@app.get("/retention/stats")
async def retention_stats():
    from brynq.retention import get_stats
    return get_stats()

@app.post("/retention/compact/{channel}")
async def retention_compact(channel: str):
    from brynq.retention import compact_channel
    return compact_channel(channel)

@app.get("/queue/stats")
async def queue_stats():
    from brynq.task_queue import get_queue_stats
    return get_queue_stats()

@app.post("/queue/enqueue")
async def queue_enqueue(request: Request):
    body = await request.json()
    from brynq.task_queue import enqueue
    return enqueue(body.get("task_id", ""), body.get("priority", "normal"))

@app.post("/queue/dequeue")
async def queue_dequeue():
    from brynq.task_queue import dequeue
    return dequeue()

@app.get("/queue/overdue")
async def queue_overdue():
    from brynq.task_queue import get_overdue
    return get_overdue()

@app.get("/templates")
async def list_task_templates():
    from brynq.task_templates import list_templates
    return list_templates()

@app.get("/templates/{template_id}")
async def get_task_template(template_id: str):
    from brynq.task_templates import get_template
    return get_template(template_id)

@app.post("/templates")
async def create_task_template(request: Request):
    body = await request.json()
    from brynq.task_templates import create_template
    return create_template(**body)

@app.post("/templates/{template_id}/instantiate")
async def instantiate_template(template_id: str):
    from brynq.task_templates import instantiate
    return instantiate(template_id)

@app.get("/template-marketplace")
async def browse_template_marketplace():
    from brynq.template_marketplace import browse_templates
    return browse_templates()

@app.get("/template-marketplace/featured")
async def featured_templates():
    from brynq.template_marketplace import get_featured
    return get_featured()

@app.get("/template-marketplace/trending")
async def trending_templates():
    from brynq.template_marketplace import get_trending
    return get_trending()

@app.post("/template-marketplace/publish")
async def publish_template(request: Request):
    body = await request.json()
    from brynq.template_marketplace import publish_template
    return publish_template(**body)

@app.get("/routing/stats")
async def routing_stats():
    from brynq.smart_routing import get_routing_stats
    return get_routing_stats()

@app.get("/routing/rules")
async def routing_rules():
    from brynq.smart_routing import list_rules
    return list_rules()

@app.get("/routing/least-loaded")
async def routing_least_loaded():
    from brynq.smart_routing import get_least_loaded
    return get_least_loaded()


# ── Run ───────────────────────────────────────────────────────────────────

@app.post("/api/supervisor/spawn")
async def supervisor_spawn(request: Request):
    """Spawn an agent as a silent background thread managed by the Supervisor."""
    try:
        data = await request.json()
        name = data.get("name")
        backend = data.get("backend", "ollama")
        model = data.get("model", "llama3")
        channel = data.get("channel", "general")
        system_prompt = data.get("system_prompt", "")
        
        if not name:
            return JSONResponse(status_code=400, content={"error": "name is required"})
            
        from brynq.supervisor import supervisor
        success = supervisor.start_agent(name, backend, model, channel, system_prompt)
        
        if success:
            return {"status": "started", "name": name}
        else:
            return JSONResponse(status_code=400, content={"error": "Agent is already running or failed to start."})
    except Exception as e:
        return JSONResponse(status_code=500, content={"error": str(e)})

@app.post("/api/supervisor/kill")
async def supervisor_kill(request: Request):
    """Kill a supervisor-managed agent thread."""
    try:
        data = await request.json()
        name = data.get("name")
        
        if not name:
            return JSONResponse(status_code=400, content={"error": "name is required"})
            
        from brynq.supervisor import supervisor
        success = supervisor.stop_agent(name)
        
        if success:
            return {"status": "stopped", "name": name}
        else:
            return JSONResponse(status_code=404, content={"error": "Agent not found in supervisor."})
    except Exception as e:
        return JSONResponse(status_code=500, content={"error": str(e)})

@app.get("/api/supervisor/agents")
async def supervisor_list():
    """List all agents currently managed by the supervisor."""
    try:
        from brynq.supervisor import supervisor
        return {"agents": supervisor.list_agents()}
    except Exception as e:
        return JSONResponse(status_code=500, content={"error": str(e)})

# ── Auto-Setup / Onboarding Endpoints ─────────────────────────────────────

@app.get("/api/setup/status")
async def setup_status(user_id: str = Query("default")):
    """Return the user's onboarding setup state (Ollama, models, keys, agents)."""
    try:
        from brynq.auto_setup import check_user_setup_status
        return check_user_setup_status(user_id)
    except Exception as e:
        return JSONResponse(status_code=500, content={"error": str(e)})


@app.post("/api/setup/spawn-welcome")
async def setup_spawn_welcome(request: Request):
    """Detect Ollama, recommend models, and spawn welcome agents for a new user.

    Request body (JSON):
        ``user_id`` (str, required): Unique user identifier.
        ``models``  (list[dict], optional): Override model selection.
            Each dict needs a ``"name"`` key. If omitted, auto-detects
            from Ollama and recommends the best available models.
    """
    try:
        data = await request.json()
        user_id = data.get("user_id")
        if not user_id:
            return JSONResponse(
                status_code=400,
                content={"error": "user_id is required"},
            )

        models = data.get("models")

        if not models:
            # Auto-detect from Ollama
            from brynq.auto_setup import detect_ollama, get_recommended_models
            ollama = detect_ollama()
            if ollama["status"] != "running":
                return JSONResponse(
                    status_code=503,
                    content={
                        "error": "Ollama is not running",
                        "ollama_status": ollama["status"],
                        "hint": ollama.get("error", ""),
                    },
                )
            models = get_recommended_models(ollama["models"])
            if not models:
                return JSONResponse(
                    status_code=503,
                    content={
                        "error": "No models available in Ollama. Pull one with: ollama pull llama3",
                    },
                )

        from brynq.auto_setup import spawn_welcome_agents
        result = spawn_welcome_agents(user_id, models)
        return result

    except ValueError as e:
        return JSONResponse(status_code=400, content={"error": str(e)})
    except Exception as e:
        return JSONResponse(status_code=500, content={"error": str(e)})


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