"""FastAPI agent endpoints — runs on every node (port 8400)."""
from __future__ import annotations

import asyncio
import logging
import uuid
from contextlib import asynccontextmanager
from datetime import datetime
from pathlib import Path

from fastapi import FastAPI, File, Form, HTTPException, UploadFile
from fastapi.responses import JSONResponse

from mesh.agent.hardware_scanner import scan_hardware
from mesh.agent.health_reporter import collect_health
from mesh.agent.job_executor import execute_job, get_active_job_count
from mesh.errors import install_error_handlers
from mesh.models import (
    HealthSnapshot, HardwareInventory, JobInfo, JobResult,
    JobStatus, JobSubmission, StatusResponse
)

logger = logging.getLogger(__name__)

# In-memory job tracking
_jobs: dict[str, JobInfo] = {}
_hardware_cache: HardwareInventory | None = None

# Set externally by start_agent.py before app starts
_node_agent = None
_mesh_db = None


def wire_agent(node_agent=None, mesh_db=None):
    """Wire in agent components before startup."""
    global _node_agent, _mesh_db
    if node_agent is not None:
        _node_agent = node_agent
    if mesh_db is not None:
        _mesh_db = mesh_db


def wire_agent_security(config, token_store=None):
    """Add auth + validation middleware to the agent app (if configured)."""
    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,
            token_store=token_store,
        )
        app.add_middleware(
            ValidationMiddleware,
            rate_limit=config.security.rate_limit,
        )
        logger.info("Agent auth middleware enabled")


@asynccontextmanager
async def lifespan(app: FastAPI):
    global _hardware_cache
    config = _node_agent.config if _node_agent is not None else None
    _hardware_cache = scan_hardware(config=config)
    logger.info("Agent API started: %s", _hardware_cache.hostname)

    if _node_agent is not None:
        asyncio.create_task(_node_agent.start())

    yield

    # Shutdown
    if _node_agent is not None:
        await _node_agent.stop()
    if _mesh_db is not None:
        _mesh_db.close()


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


@app.get("/health")
async def health():
    h = collect_health()

    # Check subsystem health
    db_ok = _mesh_db.is_healthy() if _mesh_db else True
    controller_alive = _node_agent.is_controller_alive if _node_agent else False
    last_probe = _node_agent._last_probe_time if _node_agent else 0.0

    all_ok = controller_alive or not (_node_agent and _node_agent.registered)

    data = {
        "cpu_pct": h.cpu_pct,
        "memory_used_pct": h.memory_used_pct,
        "load_1m": h.load_1m,
        "uptime_s": h.uptime_s,
        "active_jobs": get_active_job_count(),
        "db_healthy": db_ok,
        "controller_connected": controller_alive,
        "last_probe_time": last_probe,
    }

    if not all_ok and _node_agent and _node_agent.registered:
        return JSONResponse(
            status_code=503,
            content=StatusResponse(
                status="degraded",
                message="Controller connection lost",
                data=data,
            ).model_dump(mode="json"),
        )

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


@app.get("/hardware")
async def hardware() -> HardwareInventory:
    global _hardware_cache
    if _hardware_cache is None:
        config = _node_agent.config if _node_agent is not None else None
        _hardware_cache = scan_hardware(config=config)
    return _hardware_cache


@app.post("/probe/run")
async def run_probe(iterations: int = 20) -> StatusResponse:
    """Trigger a probe run in the background."""
    from mesh.agent.probe_runner import run_probes

    async def _run():
        loop = asyncio.get_event_loop()
        result = await loop.run_in_executor(
            None, lambda: run_probes(iterations=iterations)
        )
        logger.info("Probe complete: %s", result)

    asyncio.create_task(_run())
    return StatusResponse(status="ok", message="Probe started")


@app.post("/jobs/submit")
async def submit_job(submission: JobSubmission) -> JobInfo:
    """Execute a routed job on this node."""
    job_id = submission.job_id or str(uuid.uuid4())[:8]
    job = JobInfo(
        job_id=job_id,
        job_type=submission.job_type,
        command=submission.command,
        args=submission.args,
        status=JobStatus.RUNNING,
        submitted_at=datetime.utcnow(),
        started_at=datetime.utcnow(),
        priority=submission.priority,
    )
    _jobs[job_id] = job

    async def _execute():
        result = await execute_job(job)
        job.status = result.status
        job.result = result.result
        job.error = result.error
        job.completed_at = datetime.utcnow()
        # Report result back to controller
        if _node_agent and _node_agent.config.controller_url:
            try:
                from mesh.net.client import MeshClient
                async with MeshClient() as client:
                    url = f"{_node_agent.config.controller_url}/jobs/{job.job_id}/result"
                    await client.post(url, json={
                        "status": result.status.value,
                        "result": result.result,
                        "error": result.error,
                        "duration_s": result.duration_s,
                    })
            except Exception as e:
                logger.warning("Failed to report job %s result to controller: %s",
                               job.job_id, e)

    asyncio.create_task(_execute())
    return job


@app.get("/jobs/{job_id}")
async def get_job(job_id: str) -> JobInfo:
    if job_id not in _jobs:
        raise HTTPException(404, f"Job {job_id} not found")
    return _jobs[job_id]


@app.post("/tuning/apply")
async def apply_tuning(request_body: dict) -> StatusResponse:
    """Apply system tuning recommendations on this node."""
    dry_run = request_body.get("dry_run", False)
    try:
        from mesh.agent.system_optimizer import SystemOptimizer
        optimizer = SystemOptimizer(
            node_name=_hardware_cache.hostname if _hardware_cache else "unknown",
            controller_url="",
        )
        results = optimizer.apply_all(dry_run=dry_run)
        return StatusResponse(
            status="ok",
            message="Dry run complete" if dry_run else "Tuning applied",
            data=results,
        )
    except Exception as e:
        logger.error("Tuning apply failed: %s", e)
        return StatusResponse(status="error", message=str(e))


MAX_MODEL_UPLOAD_BYTES = 50 * 1024 * 1024  # 50 MB max for model uploads


@app.post("/jepa/sync")
async def jepa_sync(
    model: UploadFile = File(...),
    mean_error: str = Form("0.0"),
    training_samples: str = Form("0"),
) -> StatusResponse:
    """Receive a retrained JEPA model from the controller."""
    try:
        model_dir = Path("data")
        model_dir.mkdir(parents=True, exist_ok=True)
        model_path = model_dir / "jepa_policy_model.pt"

        # Read in chunks with size enforcement
        chunks = []
        total_size = 0
        while True:
            chunk = await model.read(1024 * 1024)  # 1MB chunks
            if not chunk:
                break
            total_size += len(chunk)
            if total_size > MAX_MODEL_UPLOAD_BYTES:
                raise HTTPException(
                    status_code=413,
                    detail=f"Model file exceeds maximum size ({MAX_MODEL_UPLOAD_BYTES // (1024*1024)} MB)",
                )
            chunks.append(chunk)

        content = b"".join(chunks)
        model_path.write_bytes(content)

        logger.info("JEPA model synced: %.1fKB, error=%s, samples=%s",
                     len(content) / 1024, mean_error, training_samples)
        return StatusResponse(
            status="ok",
            message=f"Model saved ({len(content)} bytes)",
            data={"mean_error": float(mean_error),
                  "training_samples": int(training_samples)},
        )
    except HTTPException:
        raise
    except Exception as e:
        logger.error("JEPA sync failed: %s", e)
        return StatusResponse(status="error", message="Model sync failed")


# ── Update Endpoints ────────────────────────────────────────────────────────

_agent_update_manager = None


def wire_agent_update_manager(update_manager):
    """Wire in the agent update manager."""
    global _agent_update_manager
    _agent_update_manager = update_manager


@app.post("/update/apply")
async def apply_update(
    package: UploadFile = File(...),
    sha256: str = Form(""),
) -> StatusResponse:
    """Receive and apply an update package from the controller."""
    import tempfile
    from pathlib import Path as _Path
    from mesh.updates.manager import UpdateManager

    try:
        # Save uploaded package to temp dir
        tmp_dir = _Path(tempfile.mkdtemp(prefix="mesh-agent-update-"))
        pkg_path = tmp_dir / package.filename
        content = await package.read()
        pkg_path.write_bytes(content)

        logger.info("Received update package: %s (%.1f MB)",
                     package.filename, len(content) / 1048576)

        # Get or create update manager
        mgr = _agent_update_manager
        if mgr is None:
            mgr = UpdateManager()

        # Verify SHA-256 if provided
        if sha256:
            if not mgr.verify_update(pkg_path, sha256):
                return StatusResponse(
                    status="error",
                    message="SHA-256 verification failed",
                )

        # Apply the update (restart=True to bounce the agent)
        status = mgr.apply_update(pkg_path, restart=True)

        return StatusResponse(
            status="ok" if status.value == "complete" else "error",
            message=f"Update {status.value}",
            data={"version": mgr.current_version},
        )
    except Exception as e:
        logger.error("Update apply failed: %s", e)
        return StatusResponse(status="error", message=str(e))


@app.get("/update/status")
async def update_status() -> StatusResponse:
    """Return current version and last update status."""
    from mesh.updates.manager import UpdateManager

    mgr = _agent_update_manager
    if mgr is None:
        # Default — read version from pyproject.toml or fallback
        return StatusResponse(
            status="ok",
            data={
                "version": app.version,
                "last_update": None,
            },
        )

    last_entry = None
    if mgr.history.entries:
        last_entry = mgr.history.entries[-1].model_dump(mode="json")

    return StatusResponse(
        status="ok",
        data={
            "version": mgr.current_version,
            "last_update": last_entry,
        },
    )


def create_agent_app() -> FastAPI:
    """Factory for the agent FastAPI app."""
    return app
