"""
Octopoda Cloud API Server
========================
FastAPI-based REST API for external developers to interact with Octopoda.
Runs on port 8741 (separate from the Flask dashboard on 7842).

Auto-generated docs at /docs (Swagger UI).
"""

import time
import os
import threading
from collections import OrderedDict

from fastapi import FastAPI, Depends, HTTPException, Header, Request, Query
from fastapi.middleware.cors import CORSMiddleware
from typing import Optional


# ---------------------------------------------------------------------------
# Lightweight in-process rate limiter (no external dependency)
# ---------------------------------------------------------------------------

class _RateLimiter:
    """Token-bucket rate limiter keyed by client IP."""

    def __init__(self, requests_per_minute: int = 120):
        self._rpm = requests_per_minute
        self._buckets: dict = {}  # ip -> [tokens, last_refill]
        self._lock = threading.Lock()

    def _refill(self, bucket: list):
        now = time.monotonic()
        elapsed = now - bucket[1]
        bucket[0] = min(self._rpm, bucket[0] + elapsed * (self._rpm / 60.0))
        bucket[1] = now

    def allow(self, client_ip: str) -> bool:
        with self._lock:
            if client_ip not in self._buckets:
                self._buckets[client_ip] = [self._rpm, time.monotonic()]
            bucket = self._buckets[client_ip]
            self._refill(bucket)
            if bucket[0] >= 1.0:
                bucket[0] -= 1.0
                return True
            return False


_rate_limit_rpm = int(os.environ.get("SYNRIX_RATE_LIMIT_RPM", "120"))
_rate_limiter = _RateLimiter(requests_per_minute=_rate_limit_rpm)

from synrix_runtime.api.cloud_models import (
    RegisterAgentRequest, RememberRequest, BatchRememberRequest,
    SnapshotRequest, RestoreRequest,
    SharedWriteRequest, TaskCreateRequest, TaskActionRequest, DecisionLogRequest,
    RawWriteRequest, HealthResponse, MemoryResponse, RecallResponse, SearchResponse,
    SnapshotResponse, RestoreResponse, AgentResponse, BatchMemoryResponse, ErrorResponse,
)
from synrix_runtime.api.auth import APIKeyManager
from synrix_runtime.log import get_logger

logger = get_logger("api")

# ---------------------------------------------------------------------------
# App setup
# ---------------------------------------------------------------------------

app = FastAPI(
    title="Octopoda Agent Memory API",
    version="1.0.0",
    description="Persistent Memory Kernel for AI Agents. Sub-millisecond crash recovery, shared memory bus, full audit trail.",
    docs_url="/docs",
    redoc_url="/redoc",
)

# CORS: restrict to localhost by default; set SYNRIX_CORS_ORIGINS to override
_cors_origins = os.environ.get("SYNRIX_CORS_ORIGINS", "").strip()
_allowed_origins = (
    [o.strip() for o in _cors_origins.split(",") if o.strip()]
    if _cors_origins
    else ["http://localhost:7842", "http://127.0.0.1:7842",
          "http://localhost:8741", "http://127.0.0.1:8741"]
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=_allowed_origins,
    allow_methods=["GET", "POST", "PUT", "DELETE"],
    allow_headers=["Authorization", "Content-Type"],
)


@app.middleware("http")
async def rate_limit_middleware(request: Request, call_next):
    """Enforce per-IP rate limiting on all endpoints except /health."""
    if request.url.path == "/health":
        return await call_next(request)
    client_ip = request.client.host if request.client else "unknown"
    if not _rate_limiter.allow(client_ip):
        from fastapi.responses import JSONResponse
        return JSONResponse(
            status_code=429,
            content={"detail": "Rate limit exceeded. Try again shortly."},
            headers={"Retry-After": "5"},
        )
    return await call_next(request)


# Global state (initialized by start_cloud_server())
_boot_time = time.time()
_daemon = None
_auth_manager = None
_config = None


def init_cloud_server(daemon, config):
    """Initialize the cloud server with daemon and config references."""
    global _daemon, _auth_manager, _config, _boot_time
    _daemon = daemon
    _config = config
    _boot_time = time.time()
    _auth_manager = APIKeyManager(
        backend=daemon.backend if daemon else None,
        master_key=config.api_key if config else "",
    )


# ---------------------------------------------------------------------------
# Auth dependency
# ---------------------------------------------------------------------------

async def verify_auth(authorization: Optional[str] = Header(None)):
    """Verify API key. Auth is always enforced unless SYNRIX_AUTH_DISABLED=1."""
    auth_disabled = os.environ.get("SYNRIX_AUTH_DISABLED", "").strip() == "1"
    if auth_disabled:
        return None
    if _auth_manager:
        if not authorization:
            raise HTTPException(
                status_code=401,
                detail="API key required. Pass Authorization: Bearer sk-synrix-... "
                       "(set SYNRIX_API_KEY in your environment)",
            )
        key_info = _auth_manager.verify_key(authorization)
        if not key_info:
            raise HTTPException(status_code=401, detail="Invalid API key")
        return key_info
    return None


# ---------------------------------------------------------------------------
# Helper: get or create AgentRuntime (with LRU eviction)
# ---------------------------------------------------------------------------

_agent_runtimes: OrderedDict = OrderedDict()
_MAX_CACHED_RUNTIMES = 1000


def _get_runtime(agent_id: str):
    """Get or create an AgentRuntime for the given agent (LRU cache, max 1000)."""
    if agent_id in _agent_runtimes:
        _agent_runtimes.move_to_end(agent_id)
        return _agent_runtimes[agent_id]

    # License enforcement: check agent limit before creating new runtime
    try:
        from synrix.licensing import check_agent_limit, AgentLimitError
        check_agent_limit(agent_id)
    except AgentLimitError as e:
        raise HTTPException(status_code=403, detail=str(e))
    except Exception:
        pass  # If licensing module not available, allow

    # Evict oldest if at capacity
    while len(_agent_runtimes) >= _MAX_CACHED_RUNTIMES:
        oldest_id, oldest_rt = _agent_runtimes.popitem(last=False)
        try:
            oldest_rt.shutdown()
        except Exception:
            pass
        logger.info("Evicted stale runtime: %s", oldest_id)

    from synrix_runtime.api.runtime import AgentRuntime
    rt = AgentRuntime(agent_id, agent_type="cloud")
    _agent_runtimes[agent_id] = rt
    return rt


# ---------------------------------------------------------------------------
# Health & System
# ---------------------------------------------------------------------------

@app.get("/health", response_model=HealthResponse)
async def health():
    backend_type = "unknown"
    if _daemon and hasattr(_daemon, 'backend'):
        backend_type = getattr(_daemon.backend, 'backend_type', 'unknown')
    return HealthResponse(
        status="ok",
        version="1.0.0",
        backend=backend_type,
        uptime_seconds=time.time() - _boot_time,
    )


@app.get("/v1/status")
async def system_status(auth=Depends(verify_auth)):
    if _daemon:
        return _daemon.get_system_status()
    return {"status": "no daemon", "uptime_seconds": time.time() - _boot_time}


# ---------------------------------------------------------------------------
# Agent Management (with pagination)
# ---------------------------------------------------------------------------

@app.post("/v1/agents", response_model=AgentResponse)
async def register_agent(req: RegisterAgentRequest, auth=Depends(verify_auth)):
    try:
        runtime = _get_runtime(req.agent_id)
        return AgentResponse(
            agent_id=req.agent_id,
            agent_type=req.agent_type,
            status="running",
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/v1/agents")
async def list_agents(
    offset: int = Query(default=0, ge=0),
    limit: int = Query(default=50, ge=1, le=500),
    auth=Depends(verify_auth),
):
    if _daemon:
        all_agents = _daemon.get_active_agents()
        total = len(all_agents)
        page = all_agents[offset:offset + limit]
        return {"agents": page, "count": len(page), "total": total, "offset": offset, "limit": limit}
    return {"agents": [], "count": 0, "total": 0, "offset": offset, "limit": limit}


@app.get("/v1/agents/{agent_id}")
async def get_agent(agent_id: str, auth=Depends(verify_auth)):
    if _daemon:
        state = _daemon.get_agent_state(agent_id)
        if state:
            try:
                from synrix_runtime.monitoring.metrics import MetricsCollector
                mc = MetricsCollector.get_instance()
                metrics = mc.get_agent_metrics(agent_id)
                return {
                    "agent_id": agent_id,
                    "state": state,
                    "metrics": {
                        "total_operations": metrics.total_operations,
                        "total_writes": metrics.total_writes,
                        "total_reads": metrics.total_reads,
                        "avg_write_latency_us": metrics.avg_write_latency_us,
                        "avg_read_latency_us": metrics.avg_read_latency_us,
                        "crash_count": metrics.crash_count,
                        "performance_score": metrics.performance_score,
                        "uptime_seconds": metrics.uptime_seconds,
                    },
                }
            except Exception:
                return {"agent_id": agent_id, "state": state}
    raise HTTPException(status_code=404, detail=f"Agent {agent_id} not found")


@app.delete("/v1/agents/{agent_id}")
async def deregister_agent(agent_id: str, auth=Depends(verify_auth)):
    if agent_id in _agent_runtimes:
        _agent_runtimes[agent_id].shutdown()
        del _agent_runtimes[agent_id]
    if _daemon:
        _daemon.deregister_agent(agent_id)
    return {"agent_id": agent_id, "deregistered": True}


# ---------------------------------------------------------------------------
# Memory Operations
# ---------------------------------------------------------------------------

@app.post("/v1/agents/{agent_id}/remember", response_model=MemoryResponse)
async def remember(agent_id: str, req: RememberRequest, auth=Depends(verify_auth)):
    # License enforcement: check memory limit
    try:
        from synrix.licensing import check_memory_limit, record_memory_written, MemoryLimitError
        check_memory_limit(agent_id)
    except MemoryLimitError as e:
        raise HTTPException(status_code=403, detail=str(e))
    except Exception:
        pass

    runtime = _get_runtime(agent_id)
    result = runtime.remember(req.key, req.value, tags=req.tags)

    # Track the write
    try:
        from synrix.licensing import record_memory_written
        record_memory_written(agent_id)
    except Exception:
        pass

    return MemoryResponse(
        node_id=result.node_id,
        key=req.key,
        latency_us=result.latency_us,
        timestamp=result.timestamp,
        success=result.success,
    )


@app.post("/v1/agents/{agent_id}/remember/batch", response_model=BatchMemoryResponse)
async def remember_batch(agent_id: str, req: BatchRememberRequest, auth=Depends(verify_auth)):
    runtime = _get_runtime(agent_id)
    results = []
    for item in req.items:
        # License enforcement: check memory limit per item
        try:
            from synrix.licensing import check_memory_limit, record_memory_written, MemoryLimitError
            check_memory_limit(agent_id)
        except MemoryLimitError as e:
            raise HTTPException(status_code=403, detail=str(e))
        except Exception:
            pass

        result = runtime.remember(item.key, item.value, tags=item.tags)

        try:
            record_memory_written(agent_id)
        except Exception:
            pass

        results.append({
            "key": item.key,
            "node_id": result.node_id,
            "latency_us": result.latency_us,
            "success": result.success,
        })
    return BatchMemoryResponse(agent_id=agent_id, results=results, count=len(results))


@app.get("/v1/agents/{agent_id}/recall/{key:path}", response_model=RecallResponse)
async def recall(agent_id: str, key: str, auth=Depends(verify_auth)):
    runtime = _get_runtime(agent_id)
    result = runtime.recall(key)
    return RecallResponse(
        value=result.value,
        key=key,
        latency_us=result.latency_us,
        found=result.found,
    )


@app.get("/v1/agents/{agent_id}/search", response_model=SearchResponse)
async def search(
    agent_id: str,
    prefix: str = "",
    limit: int = Query(default=50, ge=1, le=1000),
    auth=Depends(verify_auth),
):
    runtime = _get_runtime(agent_id)
    result = runtime.search(prefix, limit=limit)
    return SearchResponse(
        items=result.items,
        count=result.count,
        latency_us=result.latency_us,
    )


@app.get("/v1/agents/{agent_id}/memory")
async def list_memory(
    agent_id: str,
    offset: int = Query(default=0, ge=0),
    limit: int = Query(default=50, ge=1, le=500),
    auth=Depends(verify_auth),
):
    if _daemon and _daemon.backend:
        results = _daemon.backend.query_prefix(f"agents:{agent_id}:", limit=offset + limit)
        items = []
        for item in results:
            key = item.get("key", "")
            data = item.get("data", {})
            items.append({"key": key, "value": data.get("value", data)})
        page = items[offset:offset + limit]
        return {"agent_id": agent_id, "items": page, "count": len(page), "total": len(items), "offset": offset}
    return {"agent_id": agent_id, "items": [], "count": 0, "total": 0, "offset": offset}


# ---------------------------------------------------------------------------
# Snapshots
# ---------------------------------------------------------------------------

@app.post("/v1/agents/{agent_id}/snapshot", response_model=SnapshotResponse)
async def snapshot(agent_id: str, req: SnapshotRequest, auth=Depends(verify_auth)):
    runtime = _get_runtime(agent_id)
    result = runtime.snapshot(req.label)
    return SnapshotResponse(
        label=result.label,
        keys_captured=result.keys_captured,
        latency_us=result.latency_us,
    )


@app.post("/v1/agents/{agent_id}/restore", response_model=RestoreResponse)
async def restore(agent_id: str, req: RestoreRequest, auth=Depends(verify_auth)):
    runtime = _get_runtime(agent_id)
    result = runtime.restore(req.label)
    return RestoreResponse(
        label=result.label,
        keys_restored=result.keys_restored,
        recovery_time_us=result.recovery_time_us,
    )


# ---------------------------------------------------------------------------
# Shared Memory (with pagination)
# ---------------------------------------------------------------------------

@app.post("/v1/shared/{space}")
async def shared_write(space: str, req: SharedWriteRequest, auth=Depends(verify_auth)):
    runtime = _get_runtime(req.author_agent_id)
    result = runtime.share(req.key, req.value, space=space)
    return {
        "node_id": result.node_id,
        "key": req.key,
        "space": space,
        "latency_us": result.latency_us,
        "success": result.success,
    }


@app.get("/v1/shared/{space}/{key:path}")
async def shared_read(space: str, key: str, auth=Depends(verify_auth)):
    if _daemon and _daemon.backend:
        result = _daemon.backend.read(f"shared:{space}:{key}")
        if result:
            data = result.get("data", {})
            return {"key": key, "space": space, "value": data.get("value", data), "found": True}
    return {"key": key, "space": space, "value": None, "found": False}


@app.get("/v1/shared/{space}")
async def shared_list(
    space: str,
    offset: int = Query(default=0, ge=0),
    limit: int = Query(default=50, ge=1, le=500),
    auth=Depends(verify_auth),
):
    if _daemon and _daemon.backend:
        results = _daemon.backend.query_prefix(f"shared:{space}:", limit=offset + limit + 200)
        items = []
        for item in results:
            key = item.get("key", "").replace(f"shared:{space}:", "")
            if ":changelog:" not in key:
                data = item.get("data", {})
                items.append({"key": key, "value": data.get("value", data)})
        page = items[offset:offset + limit]
        return {"space": space, "items": page, "count": len(page), "total": len(items), "offset": offset}
    return {"space": space, "items": [], "count": 0, "total": 0, "offset": offset}


@app.get("/v1/shared")
async def shared_spaces(auth=Depends(verify_auth)):
    if _daemon and _daemon.backend:
        try:
            from synrix_runtime.api.shared_memory import SharedMemoryBus
            bus = SharedMemoryBus(_daemon.backend)
            return {"spaces": bus.list_spaces()}
        except Exception:
            pass
    return {"spaces": []}


# ---------------------------------------------------------------------------
# Audit
# ---------------------------------------------------------------------------

@app.get("/v1/agents/{agent_id}/audit")
async def agent_audit(
    agent_id: str,
    limit: int = Query(default=50, ge=1, le=500),
    auth=Depends(verify_auth),
):
    try:
        from synrix_runtime.monitoring.audit import AuditSystem
        audit = AuditSystem.get_instance()
        events = audit.replay(agent_id)
        return {"agent_id": agent_id, "events": events[:limit], "count": len(events)}
    except Exception:
        return {"agent_id": agent_id, "events": [], "count": 0}


@app.post("/v1/agents/{agent_id}/decision")
async def log_decision(agent_id: str, req: DecisionLogRequest, auth=Depends(verify_auth)):
    runtime = _get_runtime(agent_id)
    runtime.log_decision(req.decision, req.reasoning, req.context)
    return {"agent_id": agent_id, "logged": True}


# ---------------------------------------------------------------------------
# Recovery
# ---------------------------------------------------------------------------

@app.post("/v1/agents/{agent_id}/recover")
async def recover_agent(agent_id: str, auth=Depends(verify_auth)):
    if _daemon:
        result = _daemon.recover_agent(agent_id)
        return result
    raise HTTPException(status_code=503, detail="Daemon not running")


@app.get("/v1/recovery/history")
async def recovery_history(auth=Depends(verify_auth)):
    try:
        from synrix_runtime.core.recovery import RecoveryOrchestrator
        orch = RecoveryOrchestrator.get_instance()
        return {
            "history": orch.get_all_recovery_history(),
            "stats": orch.get_recovery_stats(),
        }
    except Exception:
        return {"history": [], "stats": {}}


# ---------------------------------------------------------------------------
# Metrics
# ---------------------------------------------------------------------------

@app.get("/v1/agents/{agent_id}/metrics")
async def agent_metrics(agent_id: str, auth=Depends(verify_auth)):
    try:
        from synrix_runtime.monitoring.metrics import MetricsCollector
        mc = MetricsCollector.get_instance()
        m = mc.get_agent_metrics(agent_id)
        return {
            "agent_id": agent_id,
            "total_operations": m.total_operations,
            "total_writes": m.total_writes,
            "total_reads": m.total_reads,
            "total_queries": m.total_queries,
            "avg_write_latency_us": m.avg_write_latency_us,
            "avg_read_latency_us": m.avg_read_latency_us,
            "crash_count": m.crash_count,
            "recovery_count": m.recovery_count,
            "performance_score": m.performance_score,
            "uptime_seconds": m.uptime_seconds,
            "error_rate": m.error_rate,
            "memory_node_count": m.memory_node_count,
        }
    except Exception:
        return {"agent_id": agent_id, "error": "Metrics not available"}


@app.get("/v1/metrics/system")
async def system_metrics(auth=Depends(verify_auth)):
    try:
        from synrix_runtime.monitoring.metrics import MetricsCollector
        mc = MetricsCollector.get_instance()
        m = mc.get_system_metrics()
        return {
            "total_agents": m.total_agents,
            "active_agents": m.active_agents,
            "total_operations": m.total_operations,
            "system_uptime_seconds": m.system_uptime_seconds,
            "mean_recovery_time_us": m.mean_recovery_time_us,
            "total_crashes": m.total_crashes,
            "total_recoveries": m.total_recoveries,
        }
    except Exception:
        return {"error": "Metrics not available"}


# ---------------------------------------------------------------------------
# Raw operations
# ---------------------------------------------------------------------------

@app.post("/v1/raw/write")
async def raw_write(req: RawWriteRequest, auth=Depends(verify_auth)):
    if _daemon and _daemon.backend:
        start = time.perf_counter_ns()
        node_id = _daemon.backend.write(req.key, req.value, req.metadata)
        latency_us = (time.perf_counter_ns() - start) / 1000
        return {"node_id": node_id, "key": req.key, "latency_us": round(latency_us, 1)}
    raise HTTPException(status_code=503, detail="Backend not available")


@app.get("/v1/raw/read/{key:path}")
async def raw_read(key: str, auth=Depends(verify_auth)):
    if _daemon and _daemon.backend:
        start = time.perf_counter_ns()
        result = _daemon.backend.read(key)
        latency_us = (time.perf_counter_ns() - start) / 1000
        if result:
            return {"key": key, "data": result.get("data", {}), "latency_us": round(latency_us, 1), "found": True}
        return {"key": key, "data": None, "latency_us": round(latency_us, 1), "found": False}
    raise HTTPException(status_code=503, detail="Backend not available")


@app.get("/v1/raw/query")
async def raw_query(
    prefix: str = "",
    limit: int = Query(default=100, ge=1, le=1000),
    auth=Depends(verify_auth),
):
    if _daemon and _daemon.backend:
        start = time.perf_counter_ns()
        results = _daemon.backend.query_prefix(prefix, limit=limit)
        latency_us = (time.perf_counter_ns() - start) / 1000
        return {"items": results, "count": len(results), "latency_us": round(latency_us, 1)}
    raise HTTPException(status_code=503, detail="Backend not available")


# ---------------------------------------------------------------------------
# License Info
# ---------------------------------------------------------------------------

@app.get("/v1/license")
async def license_info(auth=Depends(verify_auth)):
    """Check current license tier, agent count, and limits."""
    try:
        from synrix.licensing import get_current_claims, AgentLedger
        claims = get_current_claims()
        ledger = AgentLedger.get_instance()
        agents = ledger.get_active_agents()
        return {
            "tier": claims.tier,
            "max_agents": claims.max_agents if claims.max_agents > 0 else "unlimited",
            "max_memories_per_agent": claims.max_memories_per_agent if claims.max_memories_per_agent > 0 else "unlimited",
            "current_agents": len(agents),
            "agents": [
                {"agent_id": a, "memory_count": ledger.get_memory_count(a)}
                for a in agents
            ],
            "expires_at": claims.expires_at if claims.expires_at > 0 else None,
            "subject": claims.subject or None,
        }
    except Exception as e:
        return {"tier": "free", "error": str(e)}


# ---------------------------------------------------------------------------
# API Key Management (admin only)
# ---------------------------------------------------------------------------

@app.post("/v1/admin/keys")
async def create_api_key(tenant_id: str = "default", auth=Depends(verify_auth)):
    if _auth_manager:
        raw_key = _auth_manager.create_key(tenant_id=tenant_id)
        return {"api_key": raw_key, "tenant_id": tenant_id, "warning": "Save this key - it won't be shown again"}
    raise HTTPException(status_code=503, detail="Auth not configured")
