mx 0.1.221

use base64::{Engine as _, engine::general_purpose::STANDARD as BASE64};
use hmac::{Hmac, Mac};
use serde::{Deserialize, Serialize};
use sha2::Sha256;
use std::collections::{BTreeMap, HashMap};

use crate::knowledge::KnowledgeEntry;
use crate::store::WakeCascade;

type HmacSha256 = Hmac<Sha256>;

/// Create a signed wake ritual token: `{session_id}.{step}.{truncated_hmac[..16]}`.
///
/// `step` is a monotonic counter of chunks walked (not bloom index). The wire
/// format is unchanged from previous versions — the middle segment still parses
/// as an integer — but the semantics shift to "cumulative chunks walked" so
/// that mid-ritual bloom edits (which can change chunk counts) don't invalidate
/// previously-issued tokens. See mx#211 §2.3 / §6.
pub fn create_token(session_id: &str, step: u32) -> String {
    let payload = format!("{}.{}", session_id, step);

    let key = format!("wake-{}-ritual", session_id);
    let mut mac =
        HmacSha256::new_from_slice(key.as_bytes()).expect("HMAC can take key of any size");
    mac.update(payload.as_bytes());
    let signature = BASE64.encode(mac.finalize().into_bytes());

    format!("{}.{}", payload, &signature[..16])
}

/// Verify a wake ritual token and extract (session_id, step).
///
/// Token format: `{session_id}.{step}.{truncated_hmac[..16]}`. `step` is the
/// monotonic chunk counter; see `create_token`.
pub fn verify_token(token: &str) -> Result<(String, u32), String> {
    let parts: Vec<&str> = token.split('.').collect();
    if parts.len() != 3 {
        return Err("Invalid token format".to_string());
    }

    let session_id = parts[0];
    let step: u32 = parts[1]
        .parse()
        .map_err(|_| "Invalid step in token".to_string())?;
    let provided_sig = parts[2];

    let payload = format!("{}.{}", session_id, step);
    let key = format!("wake-{}-ritual", session_id);
    let mut mac =
        HmacSha256::new_from_slice(key.as_bytes()).expect("HMAC can take key of any size");
    mac.update(payload.as_bytes());
    let expected_sig = BASE64.encode(mac.finalize().into_bytes());

    if &expected_sig[..16] != provided_sig {
        return Err("Invalid token signature".to_string());
    }

    Ok((session_id.to_string(), step))
}

/// Where the phrase a chunk was matched against came from — authored by the
/// bloom owner, derived from the chunk's own content, or auto-generated for a
/// bloom with no authored phrases at all.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PhraseSource {
    Authored,
    Derived,
    Auto,
}

impl PhraseSource {
    pub fn as_str(self) -> &'static str {
        match self {
            PhraseSource::Authored => "authored",
            PhraseSource::Derived => "derived",
            PhraseSource::Auto => "auto",
        }
    }

    pub fn parse(s: &str) -> Option<Self> {
        match s {
            "authored" => Some(PhraseSource::Authored),
            "derived" => Some(PhraseSource::Derived),
            "auto" => Some(PhraseSource::Auto),
            _ => None,
        }
    }
}

/// Bucket counts split by phrase source. A string match against an `auto`
/// phrase that fell through to the title itself means nothing, so the split
/// keeps that visible instead of folding it into one number.
#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize, PartialEq, Eq)]
pub struct SourceCounts {
    #[serde(default)]
    pub authored: u32,
    #[serde(default)]
    pub derived: u32,
    #[serde(default)]
    pub auto: u32,
}

impl SourceCounts {
    fn bump(&mut self, source: PhraseSource) {
        match source {
            PhraseSource::Authored => self.authored += 1,
            PhraseSource::Derived => self.derived += 1,
            PhraseSource::Auto => self.auto += 1,
        }
    }

    fn total(&self) -> u32 {
        self.authored + self.derived + self.auto
    }

    fn add(&mut self, other: &SourceCounts) {
        self.authored += other.authored;
        self.derived += other.derived;
        self.auto += other.auto;
    }
}

/// Per-bloom outcome counters (1:1 with `bloom_ids`). The chunk plan itself is
/// recomputed from current content on every respond call; only the outcomes
/// accumulate here, and the final summary is their sum.
#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize)]
pub struct BloomChunkMeta {
    #[serde(default)]
    pub unhinted: SourceCounts,
    #[serde(default)]
    pub revealed: SourceCounts,
}

/// Server-side wake ritual session state.
///
/// Persisted in SurrealDB's `wake_session` table. The CLI passes a compact
/// signed token (`{session_id}.{step}.{hmac}`) between calls. State is
/// server-side; the token is just a signed reference with anti-replay.
///
/// ## Cursor invariants (Risk 4 in the design)
///
/// Two cursors compose a single position:
///
/// - `current_index` — which bloom we're on in `bloom_ids`. `0..=bloom_ids.len()`.
/// - `current_chunk_index` — which chunk within the current bloom. Always
///   advances to 0 when `current_index` advances. `0..=chunk_plan.total` for
///   the current bloom's plan.
///
/// `step` is the monotonic count of chunks walked. It is independent of
/// `current_index` / `current_chunk_index` — the latter can drift if bloom
/// content changes mid-ritual and re-chunks, but `step` always increments by
/// exactly 1 per chunk advance.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WakeSession {
    pub session_id: String,
    /// The agent that began the ritual. Carried onto every guess row, and the
    /// key every guess-log query filters on.
    pub agent: String,
    /// Wake number supplied by the caller via `--wake`. mx has no counter of
    /// its own, so this is `None` for callers that do not pass one.
    pub wake: Option<i64>,
    /// Model identifier supplied by the caller via `--model`. mx has no way to
    /// discover it.
    pub model_id: Option<String>,
    pub bloom_ids: Vec<String>,
    /// Which bloom we're on. 0-indexed; equals `bloom_ids.len()` when the
    /// ritual is complete.
    pub current_index: usize,
    /// Which chunk within the current bloom we're on. Resets to 0 when
    /// `current_index` advances. For non-chunked blooms this stays 0.
    ///
    /// Widened u8→u16 on rebase onto merged #212 — aligns with
    /// `ChunkPlan.total: u16` so large-bloom-with-low-threshold rituals
    /// (chunks > 255) address chunks correctly. Typical values remain 0-3.
    pub current_chunk_index: u16,
    /// Monotonic step counter used for token anti-replay. Ticks by 1 on every
    /// chunk advance. Survives bloom re-chunking mid-ritual.
    pub step: u32,
    pub unhinted_count: u32,
    pub revealed_count: u32,
    /// Steps walked where no guess was judged — a chunk that vanished under
    /// the cursor, or an entry deleted mid-ritual. Reported in the summary so
    /// the gap between `chunks` and the bucket totals is never a subtraction
    /// the reader has to perform.
    pub unjudged_count: u32,
    pub created_at: i64,
    /// Set by the write that walks the last step. Completed sessions are kept,
    /// not deleted, so the log can tell a finished ritual from an abandoned one
    /// and a caller whose final response was lost can still be answered.
    pub completed_at: Option<i64>,
    /// The step the last successful write started from — the token that call
    /// spent. Set by the store on every write, so a caller whose response was
    /// lost can present that token and be answered, however far the call moved
    /// `step` (a judged guess plus deleted entries stepped over is more than one).
    pub prev_step: Option<u32>,
    /// The status of the last successful write (`shown`, `bloom_missing`,
    /// `chunk_truncated`), so a resumed call that judged nothing echoes it.
    pub last_status: Option<String>,
    /// Per-bloom outcome counters (1:1 with `bloom_ids`).
    pub bloom_chunk_meta: Vec<BloomChunkMeta>,
}

impl WakeSession {
    /// Create a new session from a cascade. Chunk plans are NOT pre-computed —
    /// they are re-derived from fresh content on every respond call.
    pub fn new(
        cascade: &WakeCascade,
        agent: String,
        wake: Option<i64>,
        model_id: Option<String>,
    ) -> Self {
        let bloom_ids: Vec<String> = cascade
            .core
            .iter()
            .chain(cascade.recent.iter())
            .chain(cascade.bridges.iter())
            .map(|entry| entry.id.clone())
            .collect();
        let bloom_chunk_meta = vec![BloomChunkMeta::default(); bloom_ids.len()];

        Self {
            session_id: uuid::Uuid::new_v4().to_string(),
            agent,
            wake,
            model_id,
            bloom_ids,
            current_index: 0,
            current_chunk_index: 0,
            step: 0,
            unhinted_count: 0,
            revealed_count: 0,
            unjudged_count: 0,
            created_at: chrono::Utc::now().timestamp(),
            completed_at: None,
            prev_step: None,
            last_status: None,
            bloom_chunk_meta,
        }
    }

    /// Get current bloom ID
    pub fn current_bloom_id(&self) -> Option<&str> {
        self.bloom_ids.get(self.current_index).map(|s| s.as_str())
    }

    /// Total blooms in session
    pub fn total_blooms(&self) -> usize {
        self.bloom_ids.len()
    }

    /// Sum of the per-bloom bucket counters across the whole ritual.
    pub fn bucket_totals(&self) -> Buckets {
        let mut buckets = Buckets::default();
        for meta in &self.bloom_chunk_meta {
            buckets.unhinted.add(&meta.unhinted);
            buckets.revealed.add(&meta.revealed);
        }
        buckets
    }

    /// Current bloom position (1-indexed for display)
    pub fn current_bloom_position(&self) -> usize {
        self.current_index + 1
    }

    /// Check if ritual is complete
    pub fn is_complete(&self) -> bool {
        self.current_index >= self.bloom_ids.len()
    }

    /// Record the outcome for the current chunk and advance past it. If there
    /// are more chunks in this bloom (per `bloom_total_chunks`), tick
    /// `current_chunk_index`; otherwise advance to the next bloom and reset
    /// the chunk cursor. Always ticks `step`.
    ///
    /// Assertion-heavy by design (Risk 4): off-by-one bugs here will serve
    /// wrong content or stick the ritual.
    pub fn advance(
        &mut self,
        bloom_total_chunks: u16,
        bucket: crate::wake_guess::Bucket,
        source: PhraseSource,
    ) {
        debug_assert!(!self.is_complete(), "advance called on completed session");
        debug_assert!(
            (self.current_chunk_index as usize) < bloom_total_chunks.max(1) as usize,
            "current_chunk_index {} >= bloom_total_chunks {}",
            self.current_chunk_index,
            bloom_total_chunks
        );
        match bucket {
            crate::wake_guess::Bucket::Unhinted => self.unhinted_count += 1,
            crate::wake_guess::Bucket::Revealed => self.revealed_count += 1,
        }
        if let Some(meta) = self.bloom_chunk_meta.get_mut(self.current_index) {
            match bucket {
                crate::wake_guess::Bucket::Unhinted => meta.unhinted.bump(source),
                crate::wake_guess::Bucket::Revealed => meta.revealed.bump(source),
            }
        }
        self.step = self.step.saturating_add(1);
        self.advance_chunk_or_bloom(bloom_total_chunks);
    }

    /// Core cursor advance. Pure function of the two cursors + the chunk
    /// total. Called by `advance` above.
    fn advance_chunk_or_bloom(&mut self, bloom_total_chunks: u16) {
        let next_chunk = self.current_chunk_index.saturating_add(1);
        if (next_chunk as usize) < bloom_total_chunks.max(1) as usize {
            // More chunks in this bloom.
            self.current_chunk_index = next_chunk;
        } else {
            // Move to the next bloom; reset chunk cursor.
            self.current_index += 1;
            self.current_chunk_index = 0;
        }
        debug_assert!(
            self.current_index <= self.bloom_ids.len(),
            "current_index {} overshot bloom_ids.len() {}",
            self.current_index,
            self.bloom_ids.len()
        );
    }

    /// Handle the "content shrank mid-ritual past the cursor" case: if the
    /// recomputed chunk plan has fewer chunks than `current_chunk_index`, clamp
    /// and advance to the next bloom. Surfaced as the `chunk_truncated` status.
    ///
    /// Returns `true` if clamping occurred.
    pub fn clamp_if_chunks_shrank(&mut self, bloom_total_chunks: u16) -> bool {
        let total = bloom_total_chunks.max(1) as usize;
        if (self.current_chunk_index as usize) >= total {
            self.advance_unjudged();
            true
        } else {
            false
        }
    }

    /// Advance past a step where no guess was judged: a chunk that vanished
    /// under the cursor, or a bloom deleted mid-ritual.
    ///
    /// Ticks `step` like any other advance, so every respond retires the token
    /// it consumed and the spent one stops verifying. No bucket is recorded —
    /// nothing was guessed — so `summary.chunks` counts this step while the
    /// bucket totals do not.
    pub fn advance_unjudged(&mut self) {
        self.unjudged_count = self.unjudged_count.saturating_add(1);
        self.step = self.step.saturating_add(1);
        self.current_index += 1;
        self.current_chunk_index = 0;
    }

    /// True when the bloom the cursor points at is not in `present`.
    pub fn current_bloom_is_missing(&self, present: &HashMap<String, KnowledgeEntry>) -> bool {
        match self.current_bloom_id() {
            Some(id) => !present.contains_key(id),
            None => false,
        }
    }
}

/// Every authored phrase on an entry, in author order. `wake_phrases` takes
/// priority over the legacy single `wake_phrase`.
pub fn authored_phrases(entry: &KnowledgeEntry) -> Vec<String> {
    if !entry.wake_phrases.is_empty() {
        entry.wake_phrases.clone()
    } else if let Some(ref phrase) = entry.wake_phrase {
        vec![phrase.clone()]
    } else {
        Vec::new()
    }
}

// ============================================================================
// JSON output structures
// ============================================================================

#[derive(Debug, Serialize)]
pub struct WakeBeginResponse {
    pub status: String,
    pub session: String,
    pub prompt: BloomPrompt,
    pub progress: Progress,
    /// Per-tag counts of the entries kept out of this wake set by an excluded
    /// tag. Always present; an empty object means nothing was dropped.
    ///
    /// An UPPER BOUND, not an exact count. Each layer counts the excluded
    /// entries it walked past while filling its quota and stops counting at
    /// the quota, so an entry ranked below the wake set is never included. But
    /// an excluded entry displaces everything after it, so one that was only
    /// reached *because* an earlier exclusion pushed the window down is
    /// counted too, even though it would not have made the cut untagged.
    /// Tightening it means comparing against the untagged ordering, which for
    /// the core layer is the prefix already in hand.
    pub excluded: BTreeMap<String, usize>,
    /// The `--wake` and `--model` this ritual files every row under, echoed so
    /// a wrong number is visible before the first guess rather than after.
    pub wake: Option<i64>,
    pub model: Option<String>,
    /// Things worth a second look that did not stop the ritual, such as a wake
    /// number lower than one already logged.
    #[serde(skip_serializing_if = "Vec::is_empty")]
    pub warnings: Vec<String>,
}

#[derive(Debug, Serialize)]
pub struct WakeRespondResponse {
    /// `shown` after a guess was judged. `chunk_truncated` when the bloom
    /// shrank past the session's chunk cursor, and `bloom_missing` when the
    /// bloom was deleted mid-ritual — in both of those no guess was judged,
    /// so `bucket`, `guess` and `match` are absent and no row is logged.
    pub status: String,
    /// True when this step was already logged and the answer is the logged
    /// one: a retry after a lost response, or a second submit on one token.
    /// The first guess is the one that counts; `guess` is what was logged, not
    /// what the retry sent.
    #[serde(skip_serializing_if = "std::ops::Not::not")]
    pub replayed: bool,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub bucket: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub guess: Option<String>,
    #[serde(rename = "match", skip_serializing_if = "Option::is_none")]
    pub match_info: Option<MatchInfo>,
    /// Present on `shown` and `chunk_truncated`. Absent on `bloom_missing`,
    /// where there is no entry left to show, and on a replayed
    /// `chunk_truncated`, whose shrunk entry is not recorded.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub bloom: Option<BloomFull>,
    pub session: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub next: Option<BloomPrompt>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub progress: Option<Progress>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub summary: Option<Summary>,
    pub wake: Option<i64>,
    pub model: Option<String>,
}

/// A mechanical string-match fact. `kind` is `exact`, `close` or `none`.
#[derive(Debug, Serialize)]
pub struct MatchInfo {
    pub kind: String,
    /// Index into `bloom.phrases` of the phrase that matched. Null on `none`.
    pub phrase_index: Option<usize>,
}

#[derive(Debug, Serialize)]
pub struct WakeErrorResponse {
    pub status: String,
    pub error: String,
    pub message: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub expected_id: Option<String>,
}

/// A respond call refused with a structured payload. It is an error like any
/// other — the caller prints the JSON to stderr and exits non-zero — so a
/// wrapper that checks exit codes cannot read a mistyped id as success.
#[derive(Debug)]
pub struct WakeRejection(pub WakeErrorResponse);

impl std::fmt::Display for WakeRejection {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let json = serde_json::to_string(&self.0).map_err(|_| std::fmt::Error)?;
        f.write_str(&json)
    }
}

impl std::error::Error for WakeRejection {}

#[derive(Debug, Serialize)]
pub struct BloomPrompt {
    pub id: String,
    /// The title as shown, including any `(Part N/M)` suffix.
    pub title: String,
    pub phrase_source: String,
    /// Present only for chunked blooms. `{index: 1-based, total}`.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub chunk: Option<ChunkRef>,
}

#[derive(Debug, Serialize, Clone)]
pub struct ChunkRef {
    pub index: u16,
    pub total: u16,
    /// Present and `true` when the chunk exceeds the chunking threshold
    /// (typically an un-splittable code block). Documented limitation.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub oversized: Option<bool>,
}

#[derive(Debug, Serialize)]
pub struct BloomFull {
    pub id: String,
    pub title: String,
    /// The phrases the guess was matched against. For `derived` and `auto`
    /// sources this holds the single generated phrase.
    pub phrases: Vec<String>,
    pub phrase_source: String,
    pub content: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub chunk: Option<ChunkRef>,
}

#[derive(Debug, Serialize)]
pub struct Progress {
    /// 1-indexed count of chunks walked into.
    pub current: usize,
    /// Total chunks across the whole cascade.
    pub total: usize,
    /// 1-indexed bloom counter, and the total blooms in the cascade.
    pub bloom_current: usize,
    pub bloom_total: usize,
    /// Running bucket totals. Absent on the begin response, where nothing has
    /// been guessed yet.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub buckets: Option<BucketTotals>,
}

/// Running totals, unsplit. The per-source split is only in the final summary.
#[derive(Debug, Serialize)]
pub struct BucketTotals {
    pub unhinted: u32,
    pub revealed: u32,
}

/// Everything the model sees at the end of a ritual: what the tool did, split
/// by phrase source. No total, no ratio, no similarity value.
#[derive(Debug, Serialize)]
pub struct Summary {
    /// Every step the ritual walked, judged or not.
    pub chunks: usize,
    pub blooms: usize,
    pub buckets: Buckets,
    /// Steps where no guess was judged, because the chunk vanished under the
    /// cursor or the entry was deleted mid-ritual. `chunks` always equals the
    /// bucket totals plus this, so the difference never has to be derived —
    /// deriving it is how an "N out of M" score gets reinvented.
    pub unjudged: u32,
}

#[derive(Debug, Clone, Copy, Default, Serialize, PartialEq, Eq)]
pub struct Buckets {
    pub unhinted: SourceCounts,
    pub revealed: SourceCounts,
}

impl Buckets {
    pub fn totals(&self) -> BucketTotals {
        BucketTotals {
            unhinted: self.unhinted.total(),
            revealed: self.revealed.total(),
        }
    }
}