mod dto;
mod prompt;
mod server;
mod worker;

use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::atomic::Ordering;
use std::sync::Arc;

use anyhow::{bail, Context, Result};
use tracing::{error, info};

use crate::dto::ModelSize;
use crate::server::{AppState, ModelWorker};
use crate::worker::{WorkerConfig, WorkRequest};

// ---------------------------------------------------------------------------
// Env-var helpers
// ---------------------------------------------------------------------------

fn env_or<T: std::str::FromStr>(key: &str, default: T) -> T {
    std::env::var(key)
        .ok()
        .and_then(|v| v.parse().ok())
        .unwrap_or(default)
}

/// Parse `EMBED_MODELS` env var (comma-separated, e.g. "0.6B,2B,4B").
/// Defaults to just 0.6B.
fn parse_model_sizes() -> Vec<ModelSize> {
    let raw = std::env::var("EMBED_MODELS").unwrap_or_else(|_| "0.6B".to_string());
    raw.split(',')
        .filter_map(|s| {
            let s = s.trim();
            match s {
                "0.6B" => Some(ModelSize::S06),
                "2B" => Some(ModelSize::S2),
                "4B" => Some(ModelSize::S4),
                other => {
                    error!("Unknown model size '{other}', skipping. Valid: 0.6B, 2B, 4B");
                    None
                }
            }
        })
        .collect()
}

// ---------------------------------------------------------------------------
// Model path resolution
// ---------------------------------------------------------------------------

/// Resolve the GGUF file path for a given model size.
///
/// Priority:
///   1. Size-specific env var: `MODEL_PATH_0_6B`, `MODEL_PATH_2B`, `MODEL_PATH_4B`
///   2. Generic `MODEL_PATH` (only practical when loading a single model)
///   3. Auto-download from HuggingFace via `hf-hub`
fn resolve_model_path(size: ModelSize) -> Result<PathBuf> {
    // 1. Size-specific env var
    let size_key = match size {
        ModelSize::S06 => "MODEL_PATH_0_6B",
        ModelSize::S2 => "MODEL_PATH_2B",
        ModelSize::S4 => "MODEL_PATH_4B",
    };
    if let Ok(path) = std::env::var(size_key) {
        return Ok(PathBuf::from(path));
    }

    // 2. Generic MODEL_PATH fallback
    if let Ok(path) = std::env::var("MODEL_PATH") {
        return Ok(PathBuf::from(path));
    }

    // 3. Auto-download from HuggingFace
    info!(
        repo = size.repo_id(),
        file = size.gguf_filename(),
        "Downloading model from HuggingFace…"
    );

    let api = hf_hub::api::sync::ApiBuilder::new()
        .with_progress(true)
        .build()
        .context("unable to create HuggingFace API client")?;

    let path = api
        .model(size.repo_id().to_string())
        .get(size.gguf_filename())
        .with_context(|| format!("unable to download {}", size.gguf_filename()))?;

    Ok(path)
}

// ---------------------------------------------------------------------------
// Main
// ---------------------------------------------------------------------------

#[tokio::main]
async fn main() -> Result<()> {
    // Init tracing
    tracing_subscriber::fmt()
        .with_env_filter(
            tracing_subscriber::EnvFilter::try_from_default_env()
                .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")),
        )
        .init();

    // Parse config
    let port: u16 = env_or("EMBED_PORT", 8080);
    let max_batch: usize = env_or("EMBED_MAX_BATCH", 64);
    let max_len: usize = env_or("EMBED_MAX_LEN", 512);
    let wait_ms: u64 = env_or("EMBED_WAIT_MS", 10);

    let sizes = parse_model_sizes();
    if sizes.is_empty() {
        bail!("No valid model sizes configured. Set EMBED_MODELS=0.6B,2B,4B");
    }

    let default_size = sizes[0];
    info!(
        ?sizes,
        %default_size,
        port,
        max_batch,
        max_len,
        wait_ms,
        "Starting llama-embed-server"
    );

    // ── Init llama backend ONCE ────────────────────────────────────────
    // LlamaBackend::init() is a process-global singleton. We leak it to
    // obtain a &'static reference that all worker threads can share.
    let backend: &'static llama_cpp_2::llama_backend::LlamaBackend = {
        let b = llama_cpp_2::llama_backend::LlamaBackend::init()
            .context("LlamaBackend::init failed")?;
        Box::leak(Box::new(b))
    };

    // ── Resolve model paths and spawn workers ──────────────────────────
    let mut senders: HashMap<ModelSize, std::sync::mpsc::SyncSender<WorkRequest>> = HashMap::new();
    let mut ready_receivers = Vec::new();

    for &size in &sizes {
        let model_path = resolve_model_path(size)?;
        info!(size = %size, path = %model_path.display(), "Resolved model path");

        // Bounded channel (cap = 1024)
        let (tx, rx) = std::sync::mpsc::sync_channel::<WorkRequest>(1024);

        // Readiness oneshot
        let (ready_tx, ready_rx) = tokio::sync::oneshot::channel();

        // Worker config
        let config = WorkerConfig {
            model_path,
            max_batch,
            max_len,
            wait_ms,
        };

        // Spawn GPU worker on a dedicated OS thread
        let label = size.label().to_string();
        std::thread::Builder::new()
            .name(format!("gpu-worker-{label}"))
            .spawn(move || {
                worker::run_worker(backend, config, rx, ready_tx);
            })
            .with_context(|| format!("failed to spawn worker for {label}"))?;

        senders.insert(size, tx);
        ready_receivers.push((size, ready_rx));
    }

    // ── Await all workers to be ready ──────────────────────────────────
    let mut model_workers = HashMap::new();
    for (size, ready_rx) in ready_receivers {
        let n_embd = ready_rx
            .await
            .with_context(|| format!("{size} worker died during init"))?
            .with_context(|| format!("{size} worker init error"))?;

        info!(size = %size, n_embd, "Worker ready");

        let tx = senders.remove(&size).unwrap();
        model_workers.insert(size, ModelWorker { tx, n_embd });
    }

    // ── Build shared state and serve ───────────────────────────────────
    let state = Arc::new(AppState {
        workers: model_workers,
        ready: std::sync::atomic::AtomicBool::new(true),
        default_size,
    });

    state.ready.store(true, Ordering::Release);

    let app = server::router(state);
    let addr = format!("0.0.0.0:{port}");
    let listener = tokio::net::TcpListener::bind(&addr)
        .await
        .with_context(|| format!("failed to bind {addr}"))?;

    info!(%addr, "Listening");
    axum::serve(listener, app).await.context("axum::serve")?;

    Ok(())
}