// SPDX-License-Identifier: AGPL-3.0-or-later
// Copyright (C) 2024-2026 SuperNovae Studio <contact@supernovae.studio>

//! The candle backend — real local inference (feature `local-infer`).
//!
//! Composes the crate's candle-free control modules into the canonical candle
//! quantized generation loop (research-backed · Context7 `/huggingface/candle`
//! + the 2026-06 arXiv pass):
//!
//! ```text
//! load    gguf_file::Content::read → ModelWeights::from_gguf (Qwen3 family)
//!         tokenizers::Tokenizer::from_file (tokenizer.json — NOT in the GGUF)
//! render  ChatFamily::render(messages)            (template.rs)
//! prefill model.forward(prompt, 0)                (KV-cache lives in the model)
//! decode  model.forward([next], offset) → repeat-penalty (logits.rs) →
//!         LogitsProcessor::sample → StopController (EOS set + stop strings)
//! ```
//!
//! Device: CPU by default; Metal behind the `metal` feature. The whole module
//! runs inside the supervised child process (ADR-091) — a panic here never
//! takes down the engine. Goldens run CPU + greedy + fixed seed ONLY (Metal
//! is validated semantically, never byte-equality — backend reductions are
//! not byte-deterministic).
//!
//! `clear_kv_cache()` between requests is MANDATORY (the #1 correctness bug:
//! reusing a model across prompts without clearing desyncs cache positions →
//! garbage). The mutex below serializes requests for exactly that reason
//! (v1 = one in-flight generation per process · the documented limitation).

use std::path::Path;
use std::sync::Mutex;

use candle_core::quantized::gguf_file;
use candle_core::{Device, Tensor};
use candle_transformers::generation::{LogitsProcessor, Sampling};
use candle_transformers::models::quantized_qwen3::ModelWeights as Qwen3;

use crate::backend::BackendDyn;
use crate::error::InferLocalError;
use crate::logits::{apply_min_p, apply_repeat_penalty, apply_top_n_sigma};
use crate::protocol::{ChatRequest, ChatResponse, FinishReason, Usage};
use crate::sampling::SamplingConfig;
use crate::stop::StopController;
use crate::template::ChatFamily;

/// Default generation budget when the request does not specify `max_tokens`.
/// NOT a hard cap on its own — a request MAY ask for more. The hard bound is
/// the model context window, applied unconditionally at the decode site by
/// [`crate::budget::effective_max_tokens`] (a wire `max_tokens: u32::MAX` is
/// clamped to the slots left in the window — no OOM, no past-window desync).
const DEFAULT_MAX_TOKENS: u32 = 1024;

/// The GGUF `general.architecture` this loader's `ModelWeights` can read.
/// `load_qwen3` uses `quantized_qwen3`, which parses the `qwen3.*` metadata
/// namespace — a different arch must go through its own family loader.
const SUPPORTED_ARCH: &str = "qwen3";

/// The candle-backed local backend (Qwen3 GGUF · v1 family scope).
///
/// Owns the loaded model + tokenizer behind a `Mutex`: candle's quantized
/// model is `&mut self` on `forward` (the KV-cache lives inside), so requests
/// are serialized — v1's documented one-in-flight contract, not a perf bug.
pub struct CandleBackend {
    state: Mutex<ModelState>,
    eos_ids: Vec<u32>,
    family: ChatFamily,
    model_id: String,
    /// The device the model lives on — captured ONCE at load. Re-acquiring it
    /// per request (Metal `Device::new_metal(0)`) is a real cost + a cross-
    /// device-handle risk vs the model's tensors (review P2). `Device` is
    /// cheap to clone (an `Arc` on GPU backends · a unit on CPU).
    device: Device,
    /// The model's context window (tokens) read from the GGUF metadata —
    /// the cap a rendered prompt must not exceed (else `ContextOverflow`,
    /// instead of candle producing garbage / failing cryptically deep in
    /// the forward pass). `usize::MAX` when the metadata key is absent
    /// (degrade open · never block a load over a missing hint).
    context_window: usize,
}

struct ModelState {
    model: Qwen3,
    tokenizer: tokenizers::Tokenizer,
}

impl CandleBackend {
    /// Load a quantized Qwen3-family GGUF + its `tokenizer.json`.
    ///
    /// `device` selection: CPU always works; Metal when the `metal` feature
    /// is enabled and the host supports it (fall back to CPU otherwise).
    ///
    /// # Errors
    ///
    /// [`InferLocalError::ModelNotFound`] when either file is absent;
    /// [`InferLocalError::Backend`] for a GGUF/tokenizer parse failure.
    pub fn load_qwen3(
        gguf_path: &Path,
        tokenizer_path: &Path,
        model_id: impl Into<String>,
    ) -> Result<Self, InferLocalError> {
        let device = Self::pick_device();

        let mut file =
            std::fs::File::open(gguf_path).map_err(|_| InferLocalError::ModelNotFound {
                model: gguf_path.display().to_string(),
            })?;
        let content =
            gguf_file::Content::read(&mut file).map_err(|e| InferLocalError::Backend {
                reason: format!("GGUF parse failed: {e}"),
            })?;
        // The GGUF declares its own architecture (`general.architecture`); the
        // per-arch metadata namespace (`<arch>.context_length`, etc.) and the
        // model struct that can read it both derive from it. We VALIDATE it
        // matches what this loader's `ModelWeights` understands — handing a
        // llama/phi GGUF to the qwen3 loader otherwise fails cryptically deep
        // in from_gguf (a missing `qwen3.*` key). One canonical source, no
        // hardcoded prefix string.
        let arch = content
            .metadata
            .get("general.architecture")
            .and_then(|v| v.to_string().ok())
            .map(String::as_str)
            .ok_or_else(|| InferLocalError::Backend {
                reason: "GGUF lacks general.architecture metadata".to_owned(),
            })?
            .to_owned();
        if arch != SUPPORTED_ARCH {
            return Err(InferLocalError::Backend {
                reason: format!(
                    "this loader handles `{SUPPORTED_ARCH}` GGUFs · got `{arch}` \
                     (load via the matching family loader)"
                ),
            });
        }
        // Context window from the arch-derived key (BEFORE `content` moves).
        // Absent key → usize::MAX (degrade open · a missing hint must not
        // block the load — the overflow check simply never fires then).
        let context_window = content
            .metadata
            .get(&format!("{arch}.context_length"))
            .and_then(|v| v.to_u32().ok())
            .map_or(usize::MAX, |v| v as usize);
        let model = Qwen3::from_gguf(content, &mut file, &device).map_err(|e| {
            InferLocalError::Backend {
                reason: format!("model load failed: {e}"),
            }
        })?;

        let tokenizer = tokenizers::Tokenizer::from_file(tokenizer_path).map_err(|e| {
            InferLocalError::Backend {
                reason: format!("tokenizer load failed: {e}"),
            }
        })?;

        // Qwen3's TWO end tokens (the dual-EOS bug guard · stop.rs doc):
        // <|im_end|> ends the assistant turn · <|endoftext|> ends the text.
        let vocab = tokenizer.get_vocab(true);
        let eos_ids: Vec<u32> = ["<|im_end|>", "<|endoftext|>"]
            .iter()
            .filter_map(|t| vocab.get(*t).copied())
            .collect();
        if eos_ids.is_empty() {
            return Err(InferLocalError::Backend {
                reason: "tokenizer has neither <|im_end|> nor <|endoftext|>".to_owned(),
            });
        }

        Ok(Self {
            state: Mutex::new(ModelState { model, tokenizer }),
            eos_ids,
            family: ChatFamily::Qwen,
            model_id: model_id.into(),
            context_window,
            device,
        })
    }

    /// CPU, or Metal when compiled in AND available (silent CPU fallback —
    /// availability is a host property, not an error).
    fn pick_device() -> Device {
        #[cfg(feature = "metal")]
        {
            if let Ok(d) = Device::new_metal(0) {
                return d;
            }
        }
        Device::Cpu
    }

    /// Translate the crate's [`SamplingConfig`] into candle's `Sampling`.
    /// Greedy (the agentic default) maps to `ArgMax` — no RNG draw at all.
    fn to_candle_sampling(cfg: &SamplingConfig) -> Sampling {
        if cfg.is_greedy() {
            return Sampling::ArgMax;
        }
        match (cfg.top_k, cfg.top_p) {
            (Some(k), Some(p)) => Sampling::TopKThenTopP {
                k,
                p,
                temperature: cfg.temperature,
            },
            (Some(k), None) => Sampling::TopK {
                k,
                temperature: cfg.temperature,
            },
            (None, Some(p)) => Sampling::TopP {
                p,
                temperature: cfg.temperature,
            },
            (None, None) => Sampling::All {
                temperature: cfg.temperature,
            },
        }
    }

    /// The blocking generation entry: render → tokenize → prefill → decode
    /// loop → detokenize → response. (Runs inline in the dedicated sidecar
    /// process; the Mutex serializes generations — the v1 contract.)
    fn generate_blocking(&self, request: &ChatRequest) -> Result<ChatResponse, InferLocalError> {
        let cfg = SamplingConfig::from_request(request);
        let stop = StopController::new(self.eos_ids.iter().copied(), &request.stop);

        let mut state = self.state.lock().map_err(|_| InferLocalError::Backend {
            reason: "model mutex poisoned (a prior generation panicked)".to_owned(),
        })?;

        // Render the conversation through the family template, tokenize.
        let prompt = self.family.render(&request.messages);
        let encoding = state
            .tokenizer
            .encode(prompt.as_str(), /*add_special_tokens=*/ true)
            .map_err(|e| InferLocalError::Backend {
                reason: format!("tokenize failed: {e}"),
            })?;
        let prompt_tokens: Vec<u32> = encoding.get_ids().to_vec();
        if prompt_tokens.is_empty() {
            return Err(InferLocalError::InvalidRequest {
                reason: "empty prompt after templating".to_owned(),
            });
        }
        // Enforce the context window BEFORE the expensive prefill — a prompt
        // at/over the cap would otherwise desync rotary positions / the
        // cache and yield garbage (or fail cryptically deep in forward).
        // ≥ (not >) leaves at least one slot for a generated token.
        if prompt_tokens.len() >= self.context_window {
            return Err(InferLocalError::ContextOverflow {
                tokens: prompt_tokens.len(),
                limit: self.context_window,
            });
        }

        // Bound generation to the slots LEFT in the context window. The check
        // above guarantees prompt_tokens.len() < context_window, so this is
        // ≥ 1 — and it makes a wire `max_tokens: u32::MAX` finite (no OOM, no
        // core-pinning spin under the serialized Mutex) while preventing
        // generation past the window (rotary / KV-cache desync). The pure
        // arithmetic lives in `budget` so it is unit- + mutation-tested.
        let max_tokens = crate::budget::effective_max_tokens(
            request.max_tokens,
            DEFAULT_MAX_TOKENS,
            self.context_window,
            prompt_tokens.len(),
        );

        // MANDATORY between requests — stale cache positions produce garbage.
        state.model.clear_kv_cache();

        let mut logits_processor =
            LogitsProcessor::from_sampling(cfg.seed, Self::to_candle_sampling(&cfg));

        let (generated, finish) = decode_loop(
            &mut state,
            &self.device,
            &mut logits_processor,
            &cfg,
            &stop,
            &prompt_tokens,
            max_tokens,
        )?;

        let mut text = state
            .tokenizer
            .decode(&generated, /*skip_special_tokens=*/ true)
            .map_err(|e| InferLocalError::Backend {
                reason: format!("decode failed: {e}"),
            })?;
        if stop.hit_stop_string(&text) {
            let keep = stop.trim_stop(&text).len();
            text.truncate(keep); // trim_stop returns a prefix — truncate in place
        }

        // completion_tokens = tokens the model GENERATED this turn ("work
        // done" semantics, the standard definition). When a stop string is
        // matched its tokens are counted (the model did produce them) but
        // trimmed from the text per the OpenAI contract — a deliberate, small
        // text-vs-count asymmetry, not an over-count bug (review P2). EOS is
        // popped (a control token, never "generated content").
        let completion = u32::try_from(generated.len()).unwrap_or(u32::MAX);
        let prompt_n = u32::try_from(prompt_tokens.len()).unwrap_or(u32::MAX);
        Ok(ChatResponse::single(
            format!("local-{}", cfg.seed),
            self.model_id.clone(),
            text,
            finish,
            Usage::new(prompt_n, completion),
        ))
    }
}

/// Wrap a candle error into the crate's typed backend error — classifying
/// allocation failures as [`InferLocalError::OutOfMemory`] (candle surfaces
/// them as message text; the caller's retry policy treats OOM differently
/// from a logic error, so the type must carry the distinction).
fn backend_err(what: &str, e: impl std::fmt::Display) -> InferLocalError {
    let msg = e.to_string();
    let lower = msg.to_lowercase();
    // Phrase matches only — a bare "oom" substring false-positives on common
    // words/identifiers (bloom · room · a layer name · review P1), and a bare
    // "allocation" hits descriptive prose. Match the actual failure phrasings.
    let is_oom = lower.contains("out of memory")
        || lower.contains("cannot allocate")
        || lower.contains("failed to allocate")
        || lower.contains("allocation failed");
    if is_oom {
        return InferLocalError::OutOfMemory {
            reason: format!("{what}: {msg}"),
        };
    }
    InferLocalError::Backend {
        reason: format!("{what}: {msg}"),
    }
}

/// One sample draw, honoring the optional min-p filter (arXiv:2407.01082).
///
/// min-p is not in candle's `Sampling` enum, so it rides `sample_f` — the
/// closure candle applies to the POST-softmax probabilities right before
/// the weighted draw (our native [`apply_min_p`] zeroes the sub-floor
/// tail). candle skips the closure on the `ArgMax`/Gumbel paths, which is
/// semantically right: greedy already picks the max, truncation is moot.
fn sample_next(
    lp: &mut LogitsProcessor,
    logits: &Tensor,
    min_p: Option<f32>,
) -> Result<u32, InferLocalError> {
    let drawn = match min_p {
        Some(p) => lp.sample_f(logits, |prs| {
            apply_min_p(prs, p);
        }),
        None => lp.sample(logits),
    };
    drawn.map_err(|e| backend_err("sample", e))
}

/// Prefill + the one-token-at-a-time decode loop (the canonical candle
/// quantized path: `forward(input, offset)` with the KV-cache inside the
/// model; the offset is the ENTIRE cache contract).
fn decode_loop(
    state: &mut ModelState,
    device: &Device,
    logits_processor: &mut LogitsProcessor,
    cfg: &SamplingConfig,
    stop: &StopController,
    prompt_tokens: &[u32],
    max_tokens: usize,
) -> Result<(Vec<u32>, FinishReason), InferLocalError> {
    // ── prefill: the whole prompt at offset 0 ──
    let input = Tensor::new(prompt_tokens, device)
        .and_then(|t| t.unsqueeze(0))
        .map_err(|e| backend_err("prefill tensor", e))?;
    let logits = state
        .model
        .forward(&input, 0)
        .and_then(|l| l.squeeze(0))
        .map_err(|e| backend_err("prefill forward", e))?;
    let mut next = sample_next(logits_processor, &logits, cfg.min_p)?;

    // ── decode: one token per step, offset advancing ──
    let mut generated: Vec<u32> = vec![next];
    let has_stops = stop.has_stop_strings();
    // Per-step stop matching needs only the text TAIL as long as the longest
    // stop string — decode a bounded token suffix instead of the whole buffer
    // (O(window) per step, not O(n²) over a long generation). CONSERVATIVE
    // bound: a B-byte stop string spans AT MOST B tokens (every token decodes
    // to ≥1 byte), so `max_stop_len` tokens always contains the full stop when
    // it lands at the tail; +4 slack covers a token straddling the window edge.
    // (The earlier `/2` assumed ~2 bytes/token and could miss a stop that
    // tokenizes ~1 byte/token · review P1.) The final full-text check remains
    // a backstop, but this window now catches the stop the step it completes.
    let stop_tail_tokens = if has_stops {
        stop.max_stop_len() + 4
    } else {
        0
    };
    for index in 0..max_tokens {
        if stop.is_eos(next) {
            generated.pop(); // never include the EOS token in the text
            return Ok((generated, FinishReason::Stop));
        }
        // Stop-string check on the decoded TAIL (string-level — safe across
        // tokenizations; only when the request declared stops). A tail-decode
        // failure aborts cleanly (a tokenizer that can't decode what the model
        // produced is a real error, not a silent skip · review P2).
        if has_stops {
            let from = generated.len().saturating_sub(stop_tail_tokens);
            let tail = state
                .tokenizer
                .decode(&generated[from..], /*skip_special_tokens=*/ true)
                .map_err(|e| backend_err("tail decode", e))?;
            if stop.hit_stop_string(&tail) {
                return Ok((generated, FinishReason::Stop));
            }
        }
        if index + 1 >= max_tokens {
            break;
        }

        let input = Tensor::new(&[next], device)
            .and_then(|t| t.unsqueeze(0))
            .map_err(|e| backend_err("decode tensor", e))?;
        let mut logits = state
            .model
            .forward(&input, prompt_tokens.len() + index)
            .and_then(|l| l.squeeze(0))
            .map_err(|e| backend_err("decode forward", e))?;

        // Optional pre-softmax transforms (our native impls · logits.rs):
        // repeat penalty over the recent window, then top-nσ truncation
        // (arXiv:2411.07641 · temperature-invariant noise cut).
        if cfg.repeat_penalty > 1.0 || cfg.top_n_sigma.is_some() {
            let mut raw = logits
                .to_vec1::<f32>()
                .map_err(|e| backend_err("logits to_vec1", e))?;
            if cfg.repeat_penalty > 1.0 {
                let start = generated.len().saturating_sub(cfg.repeat_last_n);
                apply_repeat_penalty(&mut raw, cfg.repeat_penalty, &generated[start..]);
            }
            if let Some(n) = cfg.top_n_sigma {
                apply_top_n_sigma(&mut raw, n);
            }
            logits = Tensor::new(raw, device).map_err(|e| backend_err("logits rebuild", e))?;
        }

        next = sample_next(logits_processor, &logits, cfg.min_p)?;
        generated.push(next);
    }
    Ok((generated, FinishReason::Length))
}

impl BackendDyn for CandleBackend {
    async fn generate(&self, request: &ChatRequest) -> Result<ChatResponse, InferLocalError> {
        // The forward pass is CPU/GPU-bound — run it off the async thread.
        // The Mutex serializes generations (v1 single-flight contract).
        //
        // SAFETY of the block_in_place alternative rejected: spawn_blocking
        // needs 'static, and self lives in an Arc at the server layer; for
        // the v1 scaffold we run inline — the sidecar process is dedicated
        // to inference, so blocking its runtime is acceptable and documented.
        self.generate_blocking(request)
    }
}
