#![allow(
    clippy::cast_possible_wrap,
    clippy::cast_possible_truncation,
    clippy::cast_precision_loss,
    clippy::cast_sign_loss
)]

use std::num::NonZeroU32;
use std::path::PathBuf;
use std::time::{Duration, Instant};

use anyhow::Result;
use tracing::{debug, error, info, warn};

use llama_cpp_2::context::params::{LlamaContextParams, LlamaPoolingType};
use llama_cpp_2::llama_backend::LlamaBackend;
use llama_cpp_2::llama_batch::LlamaBatch;
use llama_cpp_2::model::params::LlamaModelParams;
use llama_cpp_2::model::{AddBos, LlamaModel};
use llama_cpp_2::token::LlamaToken;

use crate::dto::InputKind;
use crate::prompt;

// ---------------------------------------------------------------------------
// Public message type — sent from axum handlers to the worker thread
// ---------------------------------------------------------------------------

pub struct WorkRequest {
    pub texts: Vec<String>,
    pub input_kind: InputKind,
    pub reply: tokio::sync::oneshot::Sender<Result<Vec<Vec<f32>>>>,
}

// ---------------------------------------------------------------------------
// Worker configuration
// ---------------------------------------------------------------------------

pub struct WorkerConfig {
    pub model_path: PathBuf,
    pub max_batch: usize,
    pub max_len: usize,
    pub wait_ms: u64,
}

// ---------------------------------------------------------------------------
// GPU worker thread entry point
// ---------------------------------------------------------------------------

/// Runs on a dedicated `std::thread`. All llama model/context objects are
/// created and used exclusively on this thread (`LlamaContext` is `!Send`).
///
/// The `LlamaBackend` is initialized once on the main thread and leaked to a
/// `&'static` reference shared across all worker threads.
///
/// Sends `n_embd` (or an error) back through `ready_tx` so the main thread
/// knows the embedding dimension and can mark the server ready.
pub fn run_worker(
    backend: &'static LlamaBackend,
    config: WorkerConfig,
    rx: std::sync::mpsc::Receiver<WorkRequest>,
    ready_tx: tokio::sync::oneshot::Sender<Result<usize>>,
) {
    let WorkerConfig {
        model_path,
        max_batch,
        max_len,
        wait_ms,
    } = config;

    // ── 1. Load model ──────────────────────────────────────────────────
    let model_params = LlamaModelParams::default().with_n_gpu_layers(1000);
    info!(path = %model_path.display(), "Loading model…");

    let model = match LlamaModel::load_from_file(backend, &model_path, &model_params) {
        Ok(m) => m,
        Err(e) => {
            let err = anyhow::anyhow!("LlamaModel::load_from_file: {e}");
            error!("{err:#}");
            let _ = ready_tx.send(Err(err));
            return;
        }
    };

    // ── 2. Resolve special tokens / dimensions ─────────────────────────
    let n_embd = model.n_embd() as usize;
    info!(n_embd, "Model loaded");

    // ── 3. Create context ──────────────────────────────────────────────
    let total_ctx_tokens = max_batch * (max_len + 1); // +1 for appended EOS
    let num_cpus = std::thread::available_parallelism()
        .map(|n| n.get() as i32)
        .unwrap_or(4);

    // n_batch: max logical tokens per decode call (we want this large enough
    // to accept a single long input in one decode()).
    // n_ubatch: max PHYSICAL tokens processed at once. Smaller = less compute
    // graph memory. 2048 keeps peak workspace under ~1 GB.
    let ubatch_cap: u32 = std::env::var("EMBED_UBATCH")
        .ok()
        .and_then(|s| s.parse().ok())
        .unwrap_or(2048);
    let n_ubatch = (total_ctx_tokens as u32).min(ubatch_cap);

    let ctx_params = LlamaContextParams::default()
        .with_n_ctx(NonZeroU32::new(total_ctx_tokens as u32))
        .with_n_batch(total_ctx_tokens as u32)
        .with_n_ubatch(n_ubatch)
        .with_n_seq_max(max_batch as u32)
        .with_n_threads_batch(num_cpus)
        .with_embeddings(true)
        .with_pooling_type(LlamaPoolingType::Last); // CRITICAL for Qwen3

    let mut ctx = match model.new_context(backend, ctx_params) {
        Ok(c) => c,
        Err(e) => {
            let err = anyhow::anyhow!("new_context: {e}");
            error!("{err:#}");
            let _ = ready_tx.send(Err(err));
            return;
        }
    };

    info!(
        n_ctx = ctx.n_ctx(),
        n_batch = ctx.n_batch(),
        "Context created"
    );

    // ── 4. Signal readiness ────────────────────────────────────────────
    let _ = ready_tx.send(Ok(n_embd));

    // ── 5. Main loop — accumulate then fire ────────────────────────────
    let wait_dur = Duration::from_millis(wait_ms);

    loop {
        // Block for first request (exit on channel close)
        let first = match rx.recv() {
            Ok(req) => req,
            Err(_) => {
                info!("Channel closed, worker exiting");
                break;
            }
        };

        let mut pending: Vec<WorkRequest> = vec![first];
        let mut total_texts: usize = pending[0].texts.len();
        let deadline = Instant::now() + wait_dur;

        // Accumulate more requests until batch is full or deadline expires
        while total_texts < max_batch {
            let remaining = deadline.saturating_duration_since(Instant::now());
            if remaining.is_zero() {
                break;
            }
            match rx.recv_timeout(remaining) {
                Ok(req) => {
                    total_texts += req.texts.len();
                    pending.push(req);
                }
                Err(std::sync::mpsc::RecvTimeoutError::Timeout) => break,
                Err(std::sync::mpsc::RecvTimeoutError::Disconnected) => break,
            }
        }

        debug!(
            requests = pending.len(),
            texts = total_texts,
            "Processing batch"
        );

        process_batch(&model, &mut ctx, &mut pending, max_len, n_embd);
    }
}

// ---------------------------------------------------------------------------
// Batch processing — where correctness lives
// ---------------------------------------------------------------------------

fn process_batch(
    model: &LlamaModel,
    ctx: &mut llama_cpp_2::context::LlamaContext<'_>,
    pending: &mut Vec<WorkRequest>,
    max_len: usize,
    n_embd: usize,
) {
    // ── 1. Flatten all texts, tracking request boundaries ──────────────
    struct SeqEntry {
        request_idx: usize,
        tokens: Vec<LlamaToken>,
    }

    let mut sequences: Vec<SeqEntry> = Vec::new();

    for (req_idx, req) in pending.iter().enumerate() {
        for text in &req.texts {
            // Prompt-wrap
            let wrapped = prompt::wrap(text, req.input_kind);

            // Tokenize
            let mut tokens = match model.str_to_token(&wrapped, AddBos::Always) {
                Ok(t) => t,
                Err(e) => {
                    warn!(text_len = text.len(), "Tokenization failed: {e}");
                    // Push empty so we still have a slot; embedding will be zeros
                    Vec::new()
                }
            };

            // Truncate to max_len tokens.
            //
            // NB: str_to_token(.., AddBos::Always) calls llama_tokenize with
            // `add_special=true`, which honors the GGUF's add_bos_token /
            // add_eos_token metadata. For Qwen3-Embedding-GGUF, add_eos_token
            // is true, so EOS is already appended by the tokenizer.
            // We must NOT push another one or last-token pooling reads the
            // wrong hidden state.
            if tokens.len() > max_len {
                tokens.truncate(max_len);
            }

            sequences.push(SeqEntry {
                request_idx: req_idx,
                tokens,
            });
        }
    }

    // ── 2. Build batch ─────────────────────────────────────────────────
    let total_tokens: usize = sequences.iter().map(|s| s.tokens.len()).sum();
    let n_sequences = sequences.len();
    let seq_lens: Vec<usize> = sequences.iter().map(|s| s.tokens.len()).collect();
    info!(n_sequences, total_tokens, ?seq_lens, "Batch token stats");

    let mut batch = LlamaBatch::new(total_tokens, n_sequences as i32);
    for (seq_id, entry) in sequences.iter().enumerate() {
        if entry.tokens.is_empty() {
            continue;
        }
        if let Err(e) = batch.add_sequence(&entry.tokens, seq_id as i32, false) {
            error!(seq_id, "Failed to add sequence to batch: {e}");
        }
    }

    // ── 3. Decode ──────────────────────────────────────────────────────
    ctx.clear_kv_cache();

    if let Err(e) = ctx.decode(&mut batch) {
        error!("decode() failed: {e}");
        // Send errors to all pending requests
        for req in pending.drain(..) {
            let _ = req.reply.send(Err(anyhow::anyhow!("decode failed: {e}")));
        }
        return;
    }

    // ── 4. Extract embeddings & L2-normalize ───────────────────────────
    let mut all_embeddings: Vec<Option<Vec<f32>>> = vec![None; sequences.len()];

    for (seq_id, _entry) in sequences.iter().enumerate() {
        match ctx.embeddings_seq_ith(seq_id as i32) {
            Ok(raw) => {
                all_embeddings[seq_id] = Some(l2_normalize(raw));
            }
            Err(e) => {
                warn!(seq_id, "embeddings_seq_ith failed: {e}");
                // Leave as None; will become zero vector
            }
        }
    }

    // ── 5. Fan out — group vectors by request boundary ─────────────────
    let mut results_per_request: Vec<Vec<Vec<f32>>> = vec![Vec::new(); pending.len()];

    for (seq_id, entry) in sequences.iter().enumerate() {
        let embedding = all_embeddings[seq_id]
            .take()
            .unwrap_or_else(|| vec![0.0; n_embd]);
        results_per_request[entry.request_idx].push(embedding);
    }

    for (req_idx, req) in pending.drain(..).enumerate() {
        let vectors = std::mem::take(&mut results_per_request[req_idx]);
        let _ = req.reply.send(Ok(vectors));
    }
}

// ---------------------------------------------------------------------------
// Utility
// ---------------------------------------------------------------------------

fn l2_normalize(input: &[f32]) -> Vec<f32> {
    let magnitude = input
        .iter()
        .fold(0.0_f32, |acc, &val| val.mul_add(val, acc))
        .sqrt();

    if magnitude == 0.0 {
        return input.to_vec();
    }
    input.iter().map(|&val| val / magnitude).collect()
}