"""FastAPI controller endpoints — runs on primary node (port 8401)."""
from __future__ import annotations

import logging
import uuid
from datetime import datetime
from typing import Optional

from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import JSONResponse

from mesh.errors import install_error_handlers, NodeNotFoundError
from mesh.models import (
    AtlasSyncPayload, JEPAFeedback, JEPAStats, JobInfo, JobStatus,
    JobSubmission, NodeHeartbeat, NodeInfo, NodeRegistration,
    NodeStatus, StatusResponse, SystemTuning
)

from mesh.codex.api import router as codex_router, wire_codex

logger = logging.getLogger(__name__)

app = FastAPI(title="Mesh Controller", version="0.1.0")
install_error_handlers(app)

# Mount Codex AI agent router
app.include_router(codex_router)

# These get wired up by start_controller.py
_coordinator = None
_jepa_trainer = None
_job_router = None
_mesh_db = None
_token_store = None
_config = None
_license_client = None


def wire_controller(coordinator, jepa_trainer, job_router, mesh_db):
    """Wire in the controller components after startup."""
    global _coordinator, _jepa_trainer, _job_router, _mesh_db
    _coordinator = coordinator
    _jepa_trainer = jepa_trainer
    _job_router = job_router
    _mesh_db = mesh_db


def wire_licensing(license_client):
    """Wire in the license client for status/update endpoints."""
    global _license_client
    _license_client = license_client


def wire_security(config, token_store=None):
    """Wire in security components. Called by start_controller.py."""
    global _config, _token_store
    _config = config

    if token_store is not None:
        _token_store = token_store

    if config.security.require_auth:
        from mesh.security.auth import AuthMiddleware
        from mesh.security.validation import ValidationMiddleware
        app.add_middleware(
            AuthMiddleware,
            secret=config.security.auth_secret,
            exempt_paths=("/health", "/docs", "/openapi.json", "/nodes/register"),
            token_store=token_store,
        )
        app.add_middleware(
            ValidationMiddleware,
            rate_limit=config.security.rate_limit,
        )
        logger.info("Controller auth middleware enabled")


# ── Health Endpoint ─────────────────────────────────────────────────────────

@app.get("/health")
async def health():
    db_ok = _mesh_db.is_healthy() if _mesh_db else False
    node_count = len(_coordinator.get_all_nodes()) if _coordinator else 0
    online_count = len(_coordinator.get_online_nodes()) if _coordinator else 0
    jepa_initialized = _jepa_trainer.is_initialized if _jepa_trainer else False
    last_retrain = None
    if _jepa_trainer and _jepa_trainer._last_retrain:
        last_retrain = _jepa_trainer._last_retrain.isoformat()

    data = {
        "node_count": node_count,
        "online_nodes": online_count,
        "db_healthy": db_ok,
        "jepa_initialized": jepa_initialized,
        "last_retrain": last_retrain,
    }

    # Return 503 if critical subsystems are down
    if not db_ok:
        return JSONResponse(
            status_code=503,
            content=StatusResponse(
                status="degraded",
                message="Database unhealthy",
                data=data,
            ).model_dump(mode="json"),
        )

    return StatusResponse(status="ok", data=data)


# ── Node Endpoints ───────────────────────────────────────────────────────────

@app.post("/nodes/register")
async def register_node(reg: NodeRegistration, request: Request) -> StatusResponse:
    if _coordinator is None:
        raise HTTPException(503, "Controller not ready")

    # If auth is required, validate registration token
    if _config and _config.security.require_auth and _token_store:
        auth_header = request.headers.get("authorization", "")
        if not auth_header.startswith("Bearer "):
            raise HTTPException(401, "Registration token required")
        reg_token = auth_header[7:]
        if not _token_store.consume_registration_token(reg_token):
            raise HTTPException(401, "Invalid or already-used registration token")

    _coordinator.register_node(reg)

    # Issue an API token for the newly registered node
    response_data = {"node_id": reg.node_id}
    if _config and _token_store:
        api_token = _token_store.create_api_token(
            node_id=reg.node_id,
            role="node",
            secret=_config.security.auth_secret,
            ttl_hours=_config.security.token_ttl_hours,
        )
        response_data["api_token"] = api_token

    return StatusResponse(
        status="ok",
        message=f"Node {reg.node_id} registered",
        data=response_data,
    )


@app.post("/nodes/{node_id}/heartbeat")
async def node_heartbeat(node_id: str, hb: NodeHeartbeat) -> StatusResponse:
    if _coordinator is None:
        raise HTTPException(503, "Controller not ready")
    _coordinator.update_heartbeat(node_id, hb)
    return StatusResponse(status="ok")


@app.get("/nodes")
async def list_nodes() -> list[NodeInfo]:
    if _coordinator is None:
        return []
    return _coordinator.get_all_nodes()


@app.get("/nodes/{node_id}")
async def get_node(node_id: str) -> NodeInfo:
    if _coordinator is None:
        raise HTTPException(503, "Controller not ready")
    node = _coordinator.get_node(node_id)
    if node is None:
        raise NodeNotFoundError(node_id)
    return node


@app.delete("/nodes/{node_id}")
async def delete_node(node_id: str, archive: bool = False) -> StatusResponse:
    """Remove a node from the mesh. If archive=True, mark offline instead of deleting."""
    if _coordinator is None:
        raise HTTPException(503, "Controller not ready")
    node = _coordinator.get_node(node_id)
    if node is None:
        raise NodeNotFoundError(node_id)
    if archive:
        _coordinator.archive_node(node_id)
        return StatusResponse(status="ok", message=f"Node {node.hostname} archived")
    else:
        _coordinator.remove_node(node_id)
        return StatusResponse(status="ok", message=f"Node {node.hostname} removed")


# ── Node Capabilities ─────────────────────────────────────────────────────────

@app.get("/nodes/{node_id}/capabilities")
async def get_node_capabilities(node_id: str):
    """Return Slento CUDA, storage, and network capabilities for a node."""
    if _coordinator is None:
        raise HTTPException(503, "Controller not ready")
    node = _coordinator.get_node(node_id)
    if node is None:
        # Try by hostname
        for n in _coordinator.get_all_nodes():
            if n.hostname == node_id:
                node = n
                break
    if node is None:
        raise NodeNotFoundError(node_id)

    hw = node.hardware
    health = node.last_health
    return {
        "node_id": node.node_id,
        "hostname": node.hostname,
        "slento_cuda": {
            "available": hw.has_slento_cuda,
            "version": hw.slento_cuda_version,
            "gpu_archs": hw.slento_cuda_gpu_archs,
        },
        "gpus": [{"name": g.name, "vendor": g.vendor, "vram_mb": g.vram_mb} for g in hw.gpus],
        "storage": {
            "mounts": [m.model_dump() for m in hw.mounts],
            "scratch_paths": hw.scratch_paths,
            "mount_free_mb": health.mount_free_mb if health else {},
        },
        "network": {
            "bandwidth_mbps": health.net_bandwidth_mbps if health else 0.0,
            "lan_latency_ms": health.lan_latency_ms if health else {},
        },
        "datasets_cached": health.cached_datasets if health else [],
        "has_pytorch": hw.has_pytorch,
        "has_rocm": hw.has_rocm,
        "has_cuda": hw.has_cuda,
    }


# ── NAT-Mode Job Polling ─────────────────────────────────────────────────────
# NAT-mode agents can't receive inbound connections, so they poll this endpoint
# to pick up jobs assigned to them. The controller queues jobs for these nodes
# instead of pushing them via HTTP.

@app.get("/nodes/{node_id}/pending-jobs")
async def pending_jobs_for_node(node_id: str) -> list[dict]:
    """Return pending jobs queued for a NAT-mode node.

    Called by agents in nat_mode to pull their job assignments. Jobs are marked
    as running upon retrieval so they won't be returned again.
    """
    if _coordinator is None:
        raise HTTPException(503, "Controller not ready")
    node = _coordinator.get_node(node_id)
    if node is None:
        raise NodeNotFoundError(node_id)
    return _coordinator.get_pending_jobs_for_node(node_id)


@app.post("/jobs/{job_id}/result")
async def report_job_result(job_id: str, result: dict) -> StatusResponse:
    """Receive job completion result from an agent."""
    if _mesh_db is None:
        raise HTTPException(503, "Controller not ready")
    status = result.get("status", "completed")
    error = result.get("error")
    job_result = result.get("result")
    update = {"status": status, "error": error, "completed_at": datetime.utcnow().isoformat()}
    if job_result is not None:
        update["result"] = job_result
    _mesh_db.update_job(job_id, **update)
    logger.info("Job %s result received: %s", job_id, status)
    return StatusResponse(status="ok", message=f"Job {job_id} result recorded")


# ── Atlas Endpoints ──────────────────────────────────────────────────────────

@app.post("/atlas/sync")
async def sync_atlas(payload: AtlasSyncPayload) -> StatusResponse:
    if _coordinator is None:
        raise HTTPException(503, "Controller not ready")
    count = _coordinator.ingest_atlas_data(payload)
    return StatusResponse(
        status="ok",
        message=f"Ingested {count} data points from {payload.node_id}",
        data={"data_points": count}
    )


@app.get("/atlas/summary")
async def atlas_summary():
    from mesh.controller.federated_atlas import get_federated_summary
    if _coordinator is None:
        return {}
    return get_federated_summary(_coordinator.config.controller.atlas_dir)


# ── Job Endpoints ────────────────────────────────────────────────────────────

@app.post("/jobs/submit")
async def submit_job(submission: JobSubmission) -> JobInfo:
    if _coordinator is None:
        raise HTTPException(503, "Controller not ready")
    job = await _coordinator.submit_job(submission)
    return job


@app.get("/jobs")
async def list_jobs(status: Optional[str] = None, limit: int = 50) -> list[JobInfo]:
    if _mesh_db is None:
        return []
    rows = _mesh_db.get_jobs(status=status, limit=limit)
    return [JobInfo(**r) for r in rows]


@app.get("/jobs/{job_id}")
async def get_job(job_id: str) -> JobInfo:
    if _mesh_db is None:
        raise HTTPException(503, "Controller not ready")
    row = _mesh_db.get_job(job_id)
    if row is None:
        raise HTTPException(404, f"Job {job_id} not found")
    return JobInfo(**row)


# ── JEPA Endpoints ───────────────────────────────────────────────────────────

@app.get("/jepa/stats")
async def jepa_stats() -> JEPAStats:
    if _jepa_trainer is None:
        return JEPAStats()
    return _jepa_trainer.get_stats()


@app.post("/jepa/retrain")
async def jepa_retrain(epochs: int = 20) -> StatusResponse:
    if _jepa_trainer is None:
        raise HTTPException(503, "JEPA trainer not available")
    result = _jepa_trainer.retrain(epochs=epochs)
    return StatusResponse(status="ok", data=result)


@app.get("/jepa/schedule")
async def jepa_schedule():
    """Get current retrain schedule."""
    if _jepa_trainer is None:
        return {"interval_s": 0}
    return {"interval_s": _jepa_trainer.retrain_interval_s}


@app.post("/jepa/schedule")
async def set_jepa_schedule(interval_s: float) -> StatusResponse:
    """Update the retrain interval (in seconds)."""
    if _jepa_trainer is None:
        raise HTTPException(503, "JEPA trainer not available")
    if interval_s < 60:
        raise HTTPException(400, "Minimum interval is 60 seconds")
    _jepa_trainer.retrain_interval_s = interval_s
    return StatusResponse(status="ok", message=f"Retrain interval set to {interval_s}s")


@app.get("/jepa/training-log")
async def jepa_training_log(limit: int = 20):
    if _mesh_db is None:
        return []
    return _mesh_db.get_jepa_training_log(limit=limit)


@app.post("/jepa/feedback")
async def jepa_feedback(fb: JEPAFeedback) -> StatusResponse:
    if _jepa_trainer is None:
        raise HTTPException(503, "JEPA trainer not available")
    result = _jepa_trainer.feedback(fb)
    return StatusResponse(status="ok", data=result)


# ── System Optimization Endpoints ───────────────────────────────────────────

@app.get("/nodes/{node_id}/recommendations")
async def get_recommendations(node_id: str) -> SystemTuning:
    """Generate system tuning recommendations for a specific node."""
    if _coordinator is None:
        raise HTTPException(503, "Controller not ready")

    node = _coordinator.get_node(node_id)
    if node is None:
        # Try by hostname
        for n in _coordinator.get_all_nodes():
            if n.hostname == node_id:
                node = n
                break
    if node is None:
        raise NodeNotFoundError(node_id)

    tuning = _generate_tuning(node)
    return tuning


def _generate_tuning(node: NodeInfo) -> SystemTuning:
    """Generate tuning recommendations based on hardware + atlas data."""
    hw = node.hardware
    tuning = SystemTuning(node_id=node.node_id)

    # CPU governor — always performance for compute nodes
    tuning.cpu_governor = "performance"

    # Hugepages — scale with RAM
    ram_gb = hw.memory_total_mb / 1024
    if ram_gb >= 128:
        tuning.hugepages_2m = 4096  # 8GB hugepages
    elif ram_gb >= 64:
        tuning.hugepages_2m = 2048  # 4GB
    elif ram_gb >= 16:
        tuning.hugepages_2m = 512   # 1GB
    else:
        tuning.hugepages_2m = 128   # 256MB

    # NUMA — enable balancing for multi-socket
    tuning.numa_balancing = hw.cpu.numa_nodes > 1

    # VM tuning — lower swappiness for compute
    if ram_gb >= 64:
        tuning.swappiness = 1
        tuning.vfs_cache_pressure = 30
    else:
        tuning.swappiness = 10
        tuning.vfs_cache_pressure = 50

    # Dirty ratio — higher for big RAM (more write buffering)
    if ram_gb >= 128:
        tuning.dirty_ratio = 30
        tuning.dirty_background_ratio = 10
    else:
        tuning.dirty_ratio = 20
        tuning.dirty_background_ratio = 5

    # Scheduler (EEVDF) — disable autogroup for servers, raise util_clamp_min
    tuning.sched_autogroup = False
    if hw.cpu.cores >= 16:
        tuning.sched_util_clamp_min = 128

    # GPU nodes get more hugepages for DMA
    if hw.gpus:
        tuning.hugepages_2m = max(tuning.hugepages_2m, 2048)

    # FPGA nodes — larger net buffers for streaming data
    if hw.fpgas:
        tuning.net_rmem_max = 33554432  # 32MB
        tuning.net_wmem_max = 33554432
        tuning.net_tcp_rmem = "4096 262144 33554432"
        tuning.net_tcp_wmem = "4096 131072 33554432"

    # JEPA-informed: check if we have confidence data
    if _jepa_trainer:
        try:
            stats = _jepa_trainer.get_stats()
            if stats.confidence_by_regime:
                mem_confidence = sum(
                    v for k, v in stats.confidence_by_regime.items()
                    if "memory" in k.lower() or "bandwidth" in k.lower() or "strided" in k.lower()
                )
                if mem_confidence > 1.5:
                    tuning.hugepages_2m = max(tuning.hugepages_2m, 4096)
                    tuning.vfs_cache_pressure = 20

                compute_confidence = sum(
                    v for k, v in stats.confidence_by_regime.items()
                    if "compute" in k.lower() or "fp16" in k.lower() or "wmma" in k.lower()
                )
                if compute_confidence > 1.5:
                    tuning.sched_util_clamp_min = 256
        except Exception:
            pass

    return tuning


# ── License Endpoints ──────────────────────────────────────────────────────

@app.get("/licensing/status")
async def license_status():
    """Get current license status, tier, features, and expiry warnings."""
    if _license_client is None:
        return {
            "tier": "community",
            "valid": True,
            "max_nodes": 5,
            "features": ["dashboard", "api_access"],
            "expires_at": None,
            "warning": None,
            "license_key_set": False,
        }

    info = _license_client.info
    warning = None
    days_left = None
    if info.expires_at:
        from datetime import timezone
        now = datetime.now(timezone.utc)
        expires = info.expires_at
        if expires.tzinfo is None:
            expires = expires.replace(tzinfo=timezone.utc)
        days_left = (expires - now).days
        if not info.valid:
            warning = "License expired — running in community mode"
        elif days_left <= 14:
            warning = f"License expires in {days_left} day(s)"

    online_nodes = len(_coordinator.get_online_nodes()) if _coordinator else 0

    return {
        "tier": info.tier.value,
        "valid": info.valid,
        "customer_name": info.customer_name,
        "max_nodes": info.max_nodes,
        "current_nodes": online_nodes,
        "features": info.features,
        "issued_at": info.issued_at.isoformat() if info.issued_at else None,
        "expires_at": info.expires_at.isoformat() if info.expires_at else None,
        "days_left": days_left,
        "warning": warning,
        "license_key_set": bool(info.license_key),
    }


@app.post("/licensing/activate")
async def activate_license(request: Request) -> StatusResponse:
    """Activate or change a license key at runtime."""
    if _license_client is None:
        raise HTTPException(503, "Licensing not configured")

    body = await request.json()
    new_key = body.get("license_key", "").strip()
    if not new_key:
        raise HTTPException(400, "license_key is required")

    # Update the client's key and re-validate
    old_key = _license_client.license_key
    _license_client.license_key = new_key
    info = _license_client.validate()

    if not info.valid:
        # Revert to old key
        _license_client.license_key = old_key
        _license_client.validate()
        return StatusResponse(
            status="error",
            message="License key is invalid or expired",
        )

    logger.info("License activated: tier=%s customer=%s", info.tier.value, info.customer_name)
    return StatusResponse(
        status="ok",
        message=f"License activated — tier: {info.tier.value}, "
                f"max nodes: {info.max_nodes or 'unlimited'}",
        data={
            "tier": info.tier.value,
            "max_nodes": info.max_nodes,
            "features": info.features,
        },
    )


@app.post("/licensing/deactivate")
async def deactivate_license() -> StatusResponse:
    """Remove license key, revert to community mode."""
    if _license_client is None:
        raise HTTPException(503, "Licensing not configured")

    _license_client.license_key = ""
    _license_client.validate()
    logger.info("License deactivated — reverted to community mode")
    return StatusResponse(
        status="ok",
        message="License removed — community mode: unlimited nodes, dashboard, job routing, AMD GPU one-time probe. Continuous optimization disabled.",
    )


def create_controller_app() -> FastAPI:
    """Factory for the controller FastAPI app."""
    return app
