#!/usr/bin/env python3
"""Start the mesh controller (also starts a local agent).

Usage:
    python -m mesh.scripts.start_controller [--config mesh_config.yaml]
"""
from __future__ import annotations

import argparse
import asyncio
import logging
import sys
from pathlib import Path

import uvicorn


def main():
    parser = argparse.ArgumentParser(description="Start mesh controller")
    parser.add_argument("--config", default="mesh/mesh_config.yaml", help="Config file")
    parser.add_argument("--log-level", default=None, help="Log level")
    parser.add_argument("--no-dashboard", action="store_true", help="Skip dashboard")
    args = parser.parse_args()

    # Add project root to path
    project_root = Path(__file__).resolve().parents[2]
    sys.path.insert(0, str(project_root))

    from mesh.config import MeshConfig
    from mesh.db.mesh_db import MeshDB
    from mesh.controller.coordinator import Coordinator
    from mesh.controller.jepa_trainer import JEPATrainer
    from mesh.controller.job_router import JobRouter
    from mesh.api.controller_api import app as controller_app, wire_controller

    config = MeshConfig.from_yaml(args.config)
    if args.log_level:
        config.log_level = args.log_level

    logging.basicConfig(
        level=getattr(logging, config.log_level.upper()),
        format="%(asctime)s [%(name)s] %(levelname)s: %(message)s",
    )
    logger = logging.getLogger("mesh.controller")

    # Initialize components
    db = MeshDB(config.db_path)
    coordinator = Coordinator(config, db)
    jepa_trainer = JEPATrainer(
        atlas_dir=config.controller.atlas_dir,
        model_path=config.controller.jepa_model_path,
        db=db,
    )
    job_router = JobRouter(jepa_trainer=jepa_trainer)
    coordinator.set_job_router(job_router)
    coordinator.set_jepa_trainer(jepa_trainer)

    # Wire into FastAPI
    wire_controller(coordinator, jepa_trainer, job_router, db)

    @controller_app.on_event("startup")
    async def _start_controller():
        await coordinator.start_client()
        asyncio.create_task(coordinator.health_check_loop())

        # Periodic JEPA retrain loop
        def _get_agent_urls():
            return [n.agent_url for n in coordinator.get_online_nodes()]

        asyncio.create_task(
            jepa_trainer.retrain_loop(
                interval_s=config.controller.retrain_interval_s,
                get_agent_urls=_get_agent_urls,
            )
        )
        logger.info("Controller started on port %d (retrain every %.0fh)",
                     config.controller.api_port,
                     config.controller.retrain_interval_s / 3600)

    @controller_app.on_event("shutdown")
    async def _stop_controller():
        await coordinator.stop()

    # Mount dashboard if requested
    if not args.no_dashboard:
        try:
            from mesh.dashboard.app import create_dashboard_app
            dashboard = create_dashboard_app(
                controller_url=f"http://localhost:{config.controller.api_port}"
            )
            controller_app.mount("/dashboard", dashboard)
            logger.info("Dashboard mounted at /dashboard")
        except Exception as e:
            logger.warning("Dashboard unavailable: %s", e)

    logger.info("Starting controller on port %d", config.controller.api_port)
    uvicorn.run(controller_app, host=config.controller.host,
                port=config.controller.api_port,
                log_level=config.log_level.lower())


if __name__ == "__main__":
    main()
