// context.rs — Context window management
//
// Token counting, conversation trimming, and error classification.
// Journal entries are loaded from the memory graph store, not from
// a flat file — the parse functions are gone.

use crate::agent::api::*;
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use crate::agent::tools::working_stack;

// --- Context state types ---

/// Conversation entry — either a regular message or memory content.
/// Memory entries preserve the original message for KV cache round-tripping.
#[derive(Debug, Clone, PartialEq)]
pub enum ConversationEntry {
    /// System prompt or system-level instruction.
    System(Message),
    Message(Message),
    Memory { key: String, message: Message, score: Option<f64> },
    /// DMN heartbeat/autonomous prompt — evicted aggressively during compaction.
    Dmn(Message),
    /// Model thinking/reasoning — not sent to the API, 0 tokens for budgeting.
    Thinking(String),
    /// Debug/status log line — written to conversation log for tracing,
    /// skipped on read-back.
    Log(String),
}

/// Entry in the context window — wraps a ConversationEntry with cached metadata.
#[derive(Debug, Clone)]
pub struct ContextEntry {
    pub entry: ConversationEntry,
    /// Cached tokenization — the actual token IDs for this entry's
    /// contribution to the prompt (including chat template wrapping).
    /// Empty for Log entries.
    pub token_ids: Vec<u32>,
    /// When this entry was added to the context.
    pub timestamp: Option<DateTime<Utc>>,
}

impl ContextEntry {
    /// Create a new entry, tokenizing via the global tokenizer.
    pub fn new(entry: ConversationEntry, timestamp: Option<DateTime<Utc>>) -> Self {
        let token_ids = super::tokenizer::tokenize_conv_entry(&entry);
        Self { entry, token_ids, timestamp }
    }

    /// Token count — derived from cached token_ids length.
    pub fn tokens(&self) -> usize { self.token_ids.len() }
}

/// A named section of the context window with cached token total.
#[derive(Debug, Clone)]
pub struct ContextSection {
    pub name: String,
    /// Cached sum of entry tokens.
    tokens: usize,
    entries: Vec<ContextEntry>,
}

impl ContextSection {
    pub fn new(name: impl Into<String>) -> Self {
        Self { name: name.into(), tokens: 0, entries: Vec::new() }
    }

    pub fn entries(&self) -> &[ContextEntry] { &self.entries }
    pub fn tokens(&self) -> usize { self.tokens }
    pub fn len(&self) -> usize { self.entries.len() }
    pub fn is_empty(&self) -> bool { self.entries.is_empty() }

    /// Push a ConversationEntry, tokenizing it and updating the total.
    pub fn push_entry(&mut self, entry: ConversationEntry, timestamp: Option<DateTime<Utc>>) {
        let ce = ContextEntry::new(entry, timestamp);
        self.tokens += ce.tokens();
        self.entries.push(ce);
    }

    /// Push a pre-built ContextEntry (for restore, cloning, etc).
    pub fn push(&mut self, entry: ContextEntry) {
        self.tokens += entry.tokens();
        self.entries.push(entry);
    }

    /// Replace an entry at `index`, adjusting the token total.
    pub fn set(&mut self, index: usize, entry: ContextEntry) {
        self.tokens -= self.entries[index].tokens();
        self.tokens += entry.tokens();
        self.entries[index] = entry;
    }

    /// Remove an entry at `index`, adjusting the token total.
    pub fn del(&mut self, index: usize) -> ContextEntry {
        let removed = self.entries.remove(index);
        self.tokens -= removed.tokens();
        removed
    }

    /// Replace the message inside an entry, re-tokenizing it.
    pub fn set_message(&mut self, index: usize, msg: Message) {
        let old_tokens = self.entries[index].tokens();
        *self.entries[index].entry.message_mut() = msg;
        self.entries[index].token_ids = super::tokenizer::tokenize_conv_entry(
            &self.entries[index].entry);
        let new_tokens = self.entries[index].tokens();
        self.tokens = self.tokens - old_tokens + new_tokens;
    }

    /// Set the score on a Memory entry. No token change.
    pub fn set_score(&mut self, index: usize, score: Option<f64>) {
        if let ConversationEntry::Memory { score: s, .. } = &mut self.entries[index].entry {
            *s = score;
        }
    }

    /// Bulk replace all entries, recomputing token total.
    pub fn set_entries(&mut self, entries: Vec<ContextEntry>) {
        self.tokens = entries.iter().map(|e| e.tokens()).sum();
        self.entries = entries;
    }

    /// Dedup and trim entries to fit within context budget.
    pub fn trim(&mut self, fixed_tokens: usize) {
        let result = trim_entries(&self.entries, fixed_tokens);
        self.entries = result;
        self.tokens = self.entries.iter().map(|e| e.tokens()).sum();
    }

    /// Clear all entries.
    pub fn clear(&mut self) {
        self.entries.clear();
        self.tokens = 0;
    }
}

#[derive(Clone)]
pub struct ContextState {
    pub system: ContextSection,
    pub identity: ContextSection,
    pub journal: ContextSection,
    pub conversation: ContextSection,
    /// Working stack — separate from identity because it's managed
    /// by its own tool, not loaded from personality files.
    pub working_stack: Vec<String>,
}

impl ContextState {
    /// Total tokens across all sections.
    pub fn total_tokens(&self) -> usize {
        self.system.tokens() + self.identity.tokens()
            + self.journal.tokens() + self.conversation.tokens()
    }

    /// Budget status string for debug logging.
    pub fn format_budget(&self) -> String {
        let window = context_window();
        if window == 0 { return String::new(); }
        let used = self.total_tokens();
        let free = window.saturating_sub(used);
        let pct = |n: usize| if n == 0 { 0 } else { ((n * 100) / window).max(1) };
        format!("sys:{}% id:{}% jnl:{}% conv:{}% free:{}%",
            pct(self.system.tokens()), pct(self.identity.tokens()),
            pct(self.journal.tokens()), pct(self.conversation.tokens()),
            pct(free))
    }

    /// All sections as a slice for iteration.
    pub fn sections(&self) -> [&ContextSection; 4] {
        [&self.system, &self.identity, &self.journal, &self.conversation]
    }
}

/// Context window size in tokens (from config).
pub fn context_window() -> usize {
    crate::config::get().api_context_window
}

/// Context budget in tokens: 80% of the model's context window.
/// The remaining 20% is reserved for model output.
fn context_budget_tokens() -> usize {
    context_window() * 80 / 100
}

/// Dedup and trim conversation entries to fit within the context budget.
///
/// Phase 1: Drop duplicate memories (keep last) and DMN entries.
/// Phase 2: While over budget, drop lowest-scored memory (or if memories
///          are under 50%, drop oldest conversation entry).
fn trim_entries(entries: &[ContextEntry], fixed_tokens: usize) -> Vec<ContextEntry> {
    let max_tokens = context_budget_tokens();

    // Phase 1: dedup memories by key (keep last), drop DMN entries
    let mut seen_keys: std::collections::HashMap<&str, usize> = std::collections::HashMap::new();
    let mut drop_indices: std::collections::HashSet<usize> = std::collections::HashSet::new();

    for (i, ce) in entries.iter().enumerate() {
        if ce.entry.is_dmn() {
            drop_indices.insert(i);
        } else if let ConversationEntry::Memory { key, .. } = &ce.entry {
            if let Some(prev) = seen_keys.insert(key.as_str(), i) {
                drop_indices.insert(prev);
            }
        }
    }

    let mut result: Vec<ContextEntry> = entries.iter().enumerate()
        .filter(|(i, _)| !drop_indices.contains(i))
        .map(|(_, e)| e.clone())
        .collect();

    let entry_total = |r: &[ContextEntry]| -> usize { r.iter().map(|e| e.tokens()).sum::<usize>() };
    let mem_total = |r: &[ContextEntry]| -> usize {
        r.iter().filter(|e| e.entry.is_memory()).map(|e| e.tokens()).sum()
    };

    dbglog!("[trim] max={} fixed={} total={} entries={}",
        max_tokens, fixed_tokens, fixed_tokens + entry_total(&result), result.len());

    // Phase 2: while over budget, evict
    while fixed_tokens + entry_total(&result) > max_tokens {
        let mt = mem_total(&result);
        let ct = entry_total(&result) - mt;

        if mt > ct && let Some(i) = lowest_scored_memory(&result) {
            // If memories > 50% of entry tokens, drop lowest-scored memory
            result.remove(i);
        } else if let Some(i) = result.iter().position(|e| !e.entry.is_memory()) {
            // Otherwise drop oldest conversation entry
            result.remove(i);
        } else {
            break;
        }
    }

    // Snap to user message boundary at the start
    while let Some(first) = result.first() {
        if first.entry.message().role == Role::User { break; }
        result.remove(0);
    }

    dbglog!("[trim] result={} total={}", result.len(), fixed_tokens + entry_total(&result));
    result
}

fn lowest_scored_memory(entries: &[ContextEntry]) -> Option<usize> {
    entries.iter().enumerate()
        .filter_map(|(i, e)| match &e.entry {
            ConversationEntry::Memory { score: Some(s), .. } => Some((i, *s)),
            _ => None,
        })
        .min_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
        .map(|(i, _)| i)
}

/// Detect context window overflow errors from the API.
pub fn is_context_overflow(err: &anyhow::Error) -> bool {
    let msg = err.to_string().to_lowercase();
    msg.contains("context length")
        || msg.contains("token limit")
        || msg.contains("too many tokens")
        || msg.contains("maximum context")
        || msg.contains("prompt is too long")
        || msg.contains("request too large")
        || msg.contains("input validation error")
        || msg.contains("content length limit")
        || (msg.contains("400") && msg.contains("tokens"))
}

/// Detect model/provider errors delivered inside the SSE stream.
pub fn is_stream_error(err: &anyhow::Error) -> bool {
    err.to_string().contains("model stream error")
}

// Custom serde: serialize Memory with a "memory_key" field added to the message,
// plain messages serialize as-is. This keeps the conversation log readable.
impl Serialize for ConversationEntry {
    fn serialize<S: serde::Serializer>(&self, s: S) -> Result<S::Ok, S::Error> {
        use serde::ser::SerializeMap;
        match self {
            Self::System(m) | Self::Message(m) | Self::Dmn(m) => m.serialize(s),
            Self::Memory { key, message, score } => {
                let json = serde_json::to_value(message).map_err(serde::ser::Error::custom)?;
                let mut map = s.serialize_map(None)?;
                if let serde_json::Value::Object(obj) = json {
                    for (k, v) in obj {
                        map.serialize_entry(&k, &v)?;
                    }
                }
                map.serialize_entry("memory_key", key)?;
                if let Some(s) = score {
                    map.serialize_entry("memory_score", s)?;
                }
                map.end()
            }
            Self::Thinking(text) => {
                let mut map = s.serialize_map(Some(1))?;
                map.serialize_entry("thinking", text)?;
                map.end()
            }
            Self::Log(text) => {
                let mut map = s.serialize_map(Some(1))?;
                map.serialize_entry("log", text)?;
                map.end()
            }
        }
    }
}

impl<'de> Deserialize<'de> for ConversationEntry {
    fn deserialize<D: serde::Deserializer<'de>>(d: D) -> Result<Self, D::Error> {
        let mut json: serde_json::Value = serde_json::Value::deserialize(d)?;
        if json.get("thinking").is_some() {
            let text = json["thinking"].as_str().unwrap_or("").to_string();
            return Ok(Self::Thinking(text));
        }
        if json.get("log").is_some() {
            let text = json["log"].as_str().unwrap_or("").to_string();
            return Ok(Self::Log(text));
        }
        if let Some(key) = json.as_object_mut().and_then(|o| o.remove("memory_key")) {
            let key = key.as_str().unwrap_or("").to_string();
            let score = json.as_object_mut()
                .and_then(|o| o.remove("memory_score"))
                .and_then(|v| v.as_f64());
            let message: Message = serde_json::from_value(json).map_err(serde::de::Error::custom)?;
            Ok(Self::Memory { key, message, score })
        } else {
            let message: Message = serde_json::from_value(json).map_err(serde::de::Error::custom)?;
            Ok(Self::Message(message))
        }
    }
}

impl ConversationEntry {
    /// Get the API message for sending to the model.
    /// Panics on Log entries (which should be filtered before API calls).
    pub fn api_message(&self) -> &Message {
        match self {
            Self::System(m) | Self::Message(m) | Self::Dmn(m) => m,
            Self::Memory { message, .. } => message,
            Self::Thinking(_) | Self::Log(_) => panic!("Thinking/Log entries have no API message"),
        }
    }

    pub fn is_thinking(&self) -> bool {
        matches!(self, Self::Thinking(_))
    }

    pub fn is_memory(&self) -> bool {
        matches!(self, Self::Memory { .. })
    }

    pub fn is_dmn(&self) -> bool {
        matches!(self, Self::Dmn(_))
    }

    pub fn is_log(&self) -> bool {
        matches!(self, Self::Log(_))
    }

    /// Short description for the debug UI.
    pub fn label(&self) -> String {
        let cfg = crate::config::get();
        match self {
            Self::System(_) => "system: [system prompt]".to_string(),
            Self::Dmn(_) => "dmn: [heartbeat]".to_string(),
            Self::Thinking(text) => {
                let preview: String = text.chars().take(60).collect();
                let preview = preview.replace('\n', " ");
                if text.len() > 60 { format!("thinking: {}...", preview) }
                else { format!("thinking: {}", preview) }
            }
            Self::Log(text) => {
                let preview: String = text.chars().take(60).collect();
                format!("log: {}", preview.replace('\n', " "))
            }
            Self::Memory { key, score, .. } => {
                let role = "mem".to_string();
                match score {
                    Some(s) => format!("{}: [memory: {} score:{:.1}]", role, key, s),
                    None => format!("{}: [memory: {}]", role, key),
                }
            }
            Self::Message(m) => {
                let role = match m.role {
                    Role::Assistant => cfg.assistant_name.clone(),
                    Role::User => cfg.user_name.clone(),
                    Role::Tool => "tool".to_string(),
                    Role::System => "system".to_string(),
                };
                if let Some(tc) = &m.tool_calls {
                    let names: Vec<_> = tc.iter().map(|c| c.function.name.as_str()).collect();
                    format!("{}: [tool_call: {}]", role, names.join(", "))
                } else {
                    let text = m.content_text();
                    let preview: String = text.chars().take(60).collect();
                    let preview = preview.replace('\n', " ");
                    if text.len() > 60 { format!("{}: {}...", role, preview) }
                    else { format!("{}: {}", role, preview) }
                }
            }
        }
    }

    /// Get a reference to the inner message.
    /// Panics on Log entries.
    pub fn message(&self) -> &Message {
        match self {
            Self::System(m) | Self::Message(m) | Self::Dmn(m) => m,
            Self::Memory { message, .. } => message,
            Self::Thinking(_) | Self::Log(_) => panic!("Thinking/Log entries have no message"),
        }
    }

    /// Get a mutable reference to the inner message.
    /// Panics on Thinking/Log entries.
    pub fn message_mut(&mut self) -> &mut Message {
        match self {
            Self::System(m) | Self::Message(m) | Self::Dmn(m) => m,
            Self::Memory { message, .. } => message,
            Self::Thinking(_) | Self::Log(_) => panic!("Thinking/Log entries have no message"),
        }
    }
}

impl ContextState {
    /// Render journal entries into a single text block.
    pub fn render_journal(&self) -> String {
        if self.journal.is_empty() { return String::new(); }
        let mut text = String::from("[Earlier — from your journal]\n\n");
        for e in self.journal.entries() {
            use std::fmt::Write;
            if let Some(ts) = &e.timestamp {
                writeln!(text, "## {}\n{}\n",
                    ts.format("%Y-%m-%dT%H:%M"),
                    e.entry.message().content_text()).ok();
            } else {
                text.push_str(&e.entry.message().content_text());
                text.push_str("\n\n");
            }
        }
        text
    }

    /// Render identity files + working stack into a single user message.
    pub fn render_context_message(&self) -> String {
        let mut parts: Vec<String> = self.identity.entries().iter()
            .map(|e| e.entry.message().content_text().to_string())
            .collect();
        let instructions = std::fs::read_to_string(working_stack::instructions_path()).unwrap_or_default();
        let mut stack_section = instructions;
        if self.working_stack.is_empty() {
            stack_section.push_str("\n## Current stack\n\n(empty)\n");
        } else {
            stack_section.push_str("\n## Current stack\n\n");
            for (i, item) in self.working_stack.iter().enumerate() {
                if i == self.working_stack.len() - 1 {
                    stack_section.push_str(&format!("→ {}\n", item));
                } else {
                    stack_section.push_str(&format!("  [{}] {}\n", i, item));
                }
            }
        }
        parts.push(stack_section);
        parts.join("\n\n---\n\n")
    }
}