//! VAPORA Agent Server Binary
//! Provides HTTP server for agent coordination and health checks

use std::collections::HashMap;
use std::sync::Arc;

use anyhow::Result;
use axum::{extract::State, http::StatusCode, response::IntoResponse, routing::get, Json, Router};
use clap::Parser;
use serde_json::json;
use tokio::net::TcpListener;
use tokio::sync::mpsc;
use tracing::{error, info, warn};
use vapora_agents::{
    config::AgentConfig,
    coordinator::AgentCoordinator,
    registry::{AgentMetadata, AgentRegistry},
    runtime::executor::AgentExecutor,
};
use vapora_capabilities::CapabilityRegistry;
use vapora_llm_router::{BudgetConfig, BudgetManager, LLMRouter, LLMRouterConfig, ProviderConfig};
use vapora_shared::AgentDefinition;

#[derive(Clone)]
struct AppState {
    coordinator: Arc<AgentCoordinator>,
    #[allow(dead_code)]
    budget_manager: Option<Arc<BudgetManager>>,
    /// Path to the agent config file, used for re-loading on hot-reload.
    config_path: String,
    /// LLM router shared across executors.
    router: Option<Arc<LLMRouter>>,
    /// Capability registry for re-spawning capability executors.
    cap_registry: Arc<CapabilityRegistry>,
}

#[derive(Parser, Debug)]
#[command(
    name = "vapora-agents",
    version,
    about = "VAPORA Agent Server - Orchestration and coordination",
    long_about = "Agent server providing HTTP endpoints for agent registration, health checks, \
                  and coordination.\n\nConfiguration is loaded from:\n  - Environment variables \
                  (primary)\n  - Budget config file (default: \
                  config/agent-budgets.toml)\n\nBudget config can be overridden with \
                  --budget-config flag."
)]
struct Args {
    /// Path to budget configuration file
    #[arg(long, default_value = "config/config.ncl", env = "BUDGET_CONFIG_PATH")]
    budget_config: String,
}

#[tokio::main]
async fn main() -> Result<()> {
    let args = Args::parse();

    tracing_subscriber::fmt()
        .with_env_filter(
            tracing_subscriber::EnvFilter::from_default_env()
                .add_directive("vapora_agents=debug".parse()?),
        )
        .init();

    info!("Starting VAPORA Agent Server");

    // Load agent configuration
    let config = AgentConfig::from_env()?;
    let config_path =
        std::env::var("VAPORA_AGENT_CONFIG").unwrap_or_else(|_| "config/config.ncl".to_string());
    info!("Loaded configuration from environment");

    // Load budget configuration
    let budget_config_path = args.budget_config;
    let budget_manager = match BudgetConfig::load_or_default(&budget_config_path) {
        Ok(budget_config) => {
            if budget_config.budgets.is_empty() {
                info!(
                    "No budget configuration found at {}, running without budget enforcement",
                    budget_config_path
                );
                None
            } else {
                let manager = Arc::new(BudgetManager::new(budget_config.budgets));
                info!(
                    "Loaded budget configuration for {} roles",
                    manager.list_budgets().await.len()
                );
                Some(manager)
            }
        }
        Err(e) => {
            error!("Failed to load budget configuration: {}", e);
            return Err(e.into());
        }
    };

    // Build LLM router from config file or environment credentials
    let router = build_router_from_env();
    let router = router.map(Arc::new);

    // Initialize capability registry with built-in capability packages
    let cap_registry = Arc::new(CapabilityRegistry::with_built_ins());
    info!(
        "Capability registry initialized: {:?}",
        cap_registry.list_ids()
    );

    // Initialize agent registry and coordinator
    let max_agents_per_role = std::env::var("MAX_AGENTS_PER_ROLE")
        .ok()
        .and_then(|v| v.parse().ok())
        .unwrap_or(10);
    let registry = Arc::new(AgentRegistry::new(max_agents_per_role));
    let mut coordinator = AgentCoordinator::new(config.clone(), Arc::clone(&registry)).await?;

    if let Some(ref bm) = budget_manager {
        coordinator = coordinator.with_budget_manager(bm.clone());
        info!("Budget enforcement enabled for agent coordinator");
    }

    let coordinator = Arc::new(coordinator);

    // Spawn one executor per built-in capability, each wired to the LLM router.
    for cap_id in cap_registry.list_ids() {
        spawn_capability_executor(
            &cap_id,
            &cap_registry,
            &registry,
            &coordinator,
            router.as_ref(),
        );
    }

    // Spawn executors for config agents not covered by a capability package.
    for agent_def in &config.agents {
        if registry.get_agents_by_role(&agent_def.role).is_empty() {
            spawn_single_config_executor(agent_def, &registry, &coordinator, router.as_ref());
        }
    }

    info!("{} agents registered", registry.total_count());

    // Start coordinator
    let _coordinator_handle = {
        let coordinator = Arc::clone(&coordinator);
        tokio::spawn(async move {
            if let Err(e) = coordinator.start().await {
                error!("Coordinator error: {}", e);
            }
        })
    };

    let state = AppState {
        coordinator: Arc::clone(&coordinator),
        budget_manager,
        config_path: config_path.clone(),
        router: router.clone(),
        cap_registry: Arc::clone(&cap_registry),
    };

    // SIGHUP handler for config reload. Note: there is a brief unavailability
    // window between drain and re-registration during which assign_task returns
    // NoAvailableAgent. Learning profiles (keyed by stable_id) are preserved.
    tokio::spawn({
        let state = state.clone();
        let registry = Arc::clone(&registry);
        async move {
            let mut sighup =
                match tokio::signal::unix::signal(tokio::signal::unix::SignalKind::hangup()) {
                    Ok(s) => s,
                    Err(e) => {
                        warn!("Failed to register SIGHUP handler: {}", e);
                        return;
                    }
                };
            while sighup.recv().await.is_some() {
                handle_sighup_reload(&state, &registry).await;
            }
        }
    });

    let app = Router::new()
        .route("/health", get(health_handler))
        .route("/ready", get(readiness_handler))
        .route("/reload", axum::routing::post(reload_handler))
        .with_state(state);

    let addr = std::env::var("BIND_ADDR").unwrap_or_else(|_| "0.0.0.0:9000".to_string());
    info!("Agent server listening on {}", addr);

    let listener = TcpListener::bind(&addr).await?;
    axum::serve(listener, app).await?;

    Ok(())
}

/// Called from the SIGHUP loop. Loads a fresh config and delegates to
/// `reload_agents`. Errors are logged; the process keeps running.
async fn handle_sighup_reload(state: &AppState, registry: &Arc<AgentRegistry>) {
    info!("SIGHUP received: reloading agent configuration");
    let new_config = match AgentConfig::load(&state.config_path) {
        Ok(c) => c,
        Err(e) => {
            error!("Config parse failed during reload: {}", e);
            return;
        }
    };
    match reload_agents(
        &state.coordinator,
        registry,
        &new_config,
        &state.cap_registry,
        state.router.as_ref(),
    )
    .await
    {
        Ok(n) => info!("Reload complete: {} agents active", n),
        Err(e) => error!("Reload failed: {}", e),
    }
}

/// Drain all roles and re-spawn executors from config + capability registry.
///
/// Learning profiles (keyed by `stable_id`) are preserved in the coordinator;
/// new executor instances inherit the accumulated expertise immediately.
///
/// # Availability window
///
/// Between the drain phase and the first successful `register_agent`, callers
/// of `assign_task` for the affected roles receive `NoAvailableAgent`. This is
/// a brief window (microseconds to low milliseconds). Callers must handle this
/// error and retry. BudgetManager and LLMRouter are not reloaded; changes to
/// those require a process restart.
async fn reload_agents(
    coordinator: &Arc<AgentCoordinator>,
    registry: &Arc<AgentRegistry>,
    new_config: &AgentConfig,
    cap_registry: &Arc<CapabilityRegistry>,
    router: Option<&Arc<LLMRouter>>,
) -> anyhow::Result<usize> {
    // Drain all currently active roles. The dropped Sender causes each
    // executor's recv loop to exit after draining in-flight messages.
    for role in registry.list_roles() {
        let drained = coordinator.drain_role(&role);
        if !drained.is_empty() {
            info!("Drained {} agent(s) for role '{}'", drained.len(), role);
        }
    }

    // Re-spawn capability executors
    for cap_id in cap_registry.list_ids() {
        spawn_capability_executor(cap_id.as_str(), cap_registry, registry, coordinator, router);
    }

    // Re-spawn config agents not covered by capabilities
    for agent_def in &new_config.agents {
        if registry.get_agents_by_role(&agent_def.role).is_empty() {
            spawn_single_config_executor(agent_def, registry, coordinator, router);
        }
    }

    let total = registry.total_count();
    info!("Reload complete: {} agents active", total);
    Ok(total)
}

/// Activate a capability, register the resulting agent, and spawn its executor.
fn spawn_capability_executor(
    cap_id: &str,
    cap_registry: &CapabilityRegistry,
    registry: &Arc<AgentRegistry>,
    coordinator: &Arc<AgentCoordinator>,
    router: Option<&Arc<LLMRouter>>,
) {
    let def = match cap_registry.activate(cap_id) {
        Ok(d) => d,
        Err(e) => {
            warn!("Failed to activate capability '{}': {}", cap_id, e);
            return;
        }
    };

    let metadata = agent_metadata_from_definition(&def);
    let agent_id = metadata.id.clone();

    if let Err(e) = registry.register_agent(metadata.clone()) {
        warn!(
            "Failed to register agent for capability '{}': {}",
            cap_id, e
        );
        return;
    }

    let (tx, rx) = mpsc::channel::<vapora_agents::messages::TaskAssignment>(32);
    let executor = build_executor(metadata, rx, router);

    tokio::spawn(executor.run());
    coordinator.register_executor_channel(agent_id, tx);

    info!(
        "Spawned executor for capability '{}' (role: {}, provider: {}, model: {})",
        cap_id, def.role, def.llm_provider, def.llm_model
    );
}

fn spawn_single_config_executor(
    agent_def: &AgentDefinition,
    registry: &Arc<AgentRegistry>,
    coordinator: &Arc<AgentCoordinator>,
    router: Option<&Arc<LLMRouter>>,
) {
    let metadata = agent_metadata_from_definition(agent_def);
    let agent_id = metadata.id.clone();

    if let Err(e) = registry.register_agent(metadata.clone()) {
        warn!(
            "Failed to register config agent '{}': {}",
            agent_def.role, e
        );
        return;
    }

    let (tx, rx) = mpsc::channel::<vapora_agents::messages::TaskAssignment>(32);
    let executor = build_executor(metadata, rx, router);

    tokio::spawn(executor.run());
    coordinator.register_executor_channel(agent_id, tx);

    info!(
        "Spawned executor for config agent (role: {}, provider: {})",
        agent_def.role, agent_def.llm_provider
    );
}

fn agent_metadata_from_definition(def: &AgentDefinition) -> AgentMetadata {
    let m = AgentMetadata::new(
        def.role.clone(),
        format!("{}-agent", def.role),
        def.llm_provider.clone(),
        def.llm_model.clone(),
        def.capabilities.clone(),
    );
    match &def.system_prompt {
        Some(prompt) => m.with_system_prompt(prompt.clone()),
        None => m,
    }
}

fn build_executor(
    metadata: AgentMetadata,
    rx: mpsc::Receiver<vapora_agents::messages::TaskAssignment>,
    router: Option<&Arc<LLMRouter>>,
) -> AgentExecutor {
    let executor = AgentExecutor::new(metadata, rx);
    match router {
        Some(r) => executor.with_router(Arc::clone(r)),
        None => executor,
    }
}

/// Build an LLM router from a config file (LLM_ROUTER_CONFIG env var) or
/// from individual provider API key environment variables.
/// Returns None and logs a warning when no credentials are available.
fn build_router_from_env() -> Option<LLMRouter> {
    // Try config file path first
    if let Ok(path) = std::env::var("LLM_ROUTER_CONFIG") {
        match LLMRouterConfig::load(&path) {
            Ok(config) => match LLMRouter::new(config) {
                Ok(router) => {
                    info!("LLM router initialized from {}", path);
                    return Some(router);
                }
                Err(e) => warn!("LLM router config loaded but init failed: {}", e),
            },
            Err(e) => warn!("Failed to load LLM router config from {}: {}", path, e),
        }
    }

    // Fall back to building from individual environment variables
    let mut providers: HashMap<String, ProviderConfig> = HashMap::new();

    if let Ok(api_key) = std::env::var("ANTHROPIC_API_KEY") {
        providers.insert(
            "claude".to_string(),
            ProviderConfig {
                enabled: true,
                api_key: Some(api_key),
                url: None,
                model: "claude-sonnet-4-6".to_string(),
                max_tokens: 8192,
                temperature: 0.7,
                cost_per_1m_input: 3.0,
                cost_per_1m_output: 15.0,
            },
        );
    }

    if let Ok(api_key) = std::env::var("OPENAI_API_KEY") {
        providers.insert(
            "openai".to_string(),
            ProviderConfig {
                enabled: true,
                api_key: Some(api_key),
                url: None,
                model: "gpt-4o".to_string(),
                max_tokens: 8192,
                temperature: 0.7,
                cost_per_1m_input: 2.5,
                cost_per_1m_output: 10.0,
            },
        );
    }

    if let Ok(url) = std::env::var("OLLAMA_URL") {
        providers.insert(
            "ollama".to_string(),
            ProviderConfig {
                enabled: true,
                api_key: None,
                url: Some(url),
                model: std::env::var("OLLAMA_MODEL").unwrap_or_else(|_| "llama3.2".to_string()),
                max_tokens: 4096,
                temperature: 0.7,
                cost_per_1m_input: 0.0,
                cost_per_1m_output: 0.0,
            },
        );
    }

    if providers.is_empty() {
        info!("No LLM provider credentials found — executors will use stub responses");
        return None;
    }

    let default_provider = if providers.contains_key("claude") {
        "claude".to_string()
    } else {
        providers.keys().next().unwrap().clone()
    };

    let config = LLMRouterConfig {
        routing: vapora_llm_router::config::RoutingConfig {
            default_provider,
            cost_tracking_enabled: true,
            fallback_enabled: true,
        },
        providers,
        routing_rules: vec![],
    };

    match LLMRouter::new(config) {
        Ok(router) => {
            info!("LLM router initialized from environment credentials");
            Some(router)
        }
        Err(e) => {
            warn!("Failed to initialize LLM router from env vars: {}", e);
            None
        }
    }
}

async fn health_handler() -> Json<serde_json::Value> {
    Json(json!({
        "status": "healthy",
        "service": "vapora-agents",
        "version": env!("CARGO_PKG_VERSION")
    }))
}

async fn readiness_handler(State(state): State<AppState>) -> Json<serde_json::Value> {
    let is_ready = state.coordinator.is_ready().await;
    Json(json!({
        "ready": is_ready,
        "service": "vapora-agents",
        "agents": state.coordinator.get_agent_count().await
    }))
}

async fn reload_handler(State(state): State<AppState>) -> impl IntoResponse {
    let new_config = match AgentConfig::load(&state.config_path) {
        Ok(c) => c,
        Err(e) => {
            return (
                StatusCode::INTERNAL_SERVER_ERROR,
                Json(json!({"error": e.to_string()})),
            )
                .into_response();
        }
    };

    let registry = state.coordinator.registry_arc();
    match reload_agents(
        &state.coordinator,
        &registry,
        &new_config,
        &state.cap_registry,
        state.router.as_ref(),
    )
    .await
    {
        Ok(n) => (StatusCode::OK, Json(json!({"reloaded": true, "agents": n}))).into_response(),
        Err(e) => (
            StatusCode::INTERNAL_SERVER_ERROR,
            Json(json!({"error": e.to_string()})),
        )
            .into_response(),
    }
}
