//! One pet line or diary page: what to ask the brain, how lookups continue
//! the same line, and how the result is filtered. The runtimes only generate.
use crate::{
    types::{DiaryPage, WorldState},
    world,
};
use serde::{Deserialize, Serialize};

/// A request to the brain: continue `header`, the newest `lines` that fit and
/// `prefix`. Prompts end at `Name:` with no trailing space. Sampling keeps the
/// `top_k` likeliest tokens, then the smallest set holding `top_p` of their
/// probability; temperature 0 is greedy (used only by checks).
#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
#[serde(rename_all = "camelCase")]
pub struct Generation {
    pub header: String,
    pub lines: Vec<String>,
    pub prefix: String,
    pub mode: String,
    pub max_new: usize,
    pub temperature: f32,
    pub top_k: usize,
    pub top_p: f32,
    pub ctx: usize,
}
pub const TEMPERATURE: f32 = 0.85;
pub const TOP_K: usize = 40;
pub const TOP_P: f32 = 0.92;
/// Lookups per line before the line must end.
const MAX_LOOKUPS: u32 = 3;
/// Filtered lines are generated again this many times in all.
const TRIES: u32 = 3;

/// The pet's line, with up to three lookup round-trips. Idle turns see a
/// shorter history: prefill is most of their cost.
#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct ChatTurn {
    pub kind: String,
    world: WorldState,
    header: String,
    lines: Vec<String>,
    text: String,
    hop: u32,
    tries: u32,
    body: Option<String>,
}
impl ChatTurn {
    pub fn new(world: WorldState, lines: Vec<String>, kind: &str) -> Self {
        Self {
            kind: if kind == "idle" { "idle" } else { "react" }.into(),
            header: world::header("chat", &world),
            world,
            lines,
            text: String::new(),
            hop: 0,
            tries: 0,
            body: None,
        }
    }
    pub fn request(&self) -> Generation {
        let idle = self.kind == "idle";
        Generation {
            header: self.header.clone(),
            lines: self.lines.clone(),
            prefix: format!("{}:{}", self.world.pet, self.text),
            mode: "line".into(),
            max_new: if idle { 40 } else { 64 },
            temperature: TEMPERATURE,
            top_k: TOP_K,
            top_p: TOP_P,
            ctx: if idle { 512 } else { 1024 },
        }
    }
    /// Adds a generated piece that stopped for `reason`. Returns true when the
    /// line needs another generation: after a lookup (its result is appended
    /// here) or after a filtered line.
    pub fn feed(&mut self, text: &str, reason: &str) -> bool {
        if self.body.is_some() {
            return false;
        }
        self.text += text;
        if reason == "lookup" && self.hop < MAX_LOOKUPS {
            if let Some(result) = world::lookup_tail(&self.text)
                .and_then(|(name, arg)| world::tool_result(&self.world, &name, &arg))
            {
                self.text += &format!("{{{result}}}");
                self.hop += 1;
                return true;
            }
        }
        if world::not_cute(&self.text) {
            self.tries += 1;
            self.text.clear();
            self.hop = 0;
            if self.tries < TRIES {
                return true;
            }
            self.body = Some(String::new());
            return false;
        }
        self.body = Some(world::compact(&self.text));
        false
    }
    /// The line so far, with lookup results, for live display.
    pub fn text(&self) -> &str {
        &self.text
    }
    /// The finished line, empty when every try was filtered.
    pub fn body(&self) -> Option<&str> {
        self.body.as_deref()
    }
}

/// A diary page for one day, written from that day's log.
#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct DiaryTurn {
    pub date: String,
    pub day: i64,
    pub weather: String,
    pub stage: String,
    /// Debug rewrites replace today's page.
    pub replace: bool,
    header: String,
}
pub const DIARY_START: &str = "dear diary,";
impl DiaryTurn {
    pub fn new(world: &WorldState, date: String, replace: bool) -> Self {
        Self {
            date,
            day: world.day,
            weather: world.weather.clone(),
            stage: world.stage.clone(),
            replace,
            header: world::header("diary", world),
        }
    }
    pub fn request(&self) -> Generation {
        Generation {
            header: self.header.clone(),
            lines: vec![],
            prefix: DIARY_START.into(),
            mode: "diary".into(),
            max_new: 220,
            temperature: TEMPERATURE,
            top_k: TOP_K,
            top_p: TOP_P,
            ctx: 1024,
        }
    }
    /// The page from the brain's continuation: up to three known doodle
    /// stickers, text without tags. A page with unsafe text is dropped whole
    /// rather than publishing its fragments; the day stays pending.
    pub fn finish(&self, generated: &str) -> Option<DiaryPage> {
        let text = format!("{DIARY_START}{generated}");
        if world::not_cute(&text) {
            return None;
        }
        let doodles = world::tags(&text)
            .into_iter()
            .filter(|t| t.name == "doodle")
            .filter_map(|t| t.arg)
            .filter(|d| world::DOODLES.contains(&d.as_str()))
            .take(3)
            .collect();
        Some(DiaryPage {
            date: self.date.clone(),
            day: self.day,
            weather: self.weather.clone(),
            stage: self.stage.clone(),
            text: world::clean_page(&text),
            doodles,
        })
    }
}