ccd-cli 1.0.0-beta.5

#![allow(dead_code)]

use std::collections::HashMap;
use std::fmt;

use serde_json::Value as JsonValue;

pub const DIALECT: &str = "CCD1";
pub const FRAME_BYTES_MAX: usize = 2048;
pub const HEADER_BYTES_MAX: usize = 160;
pub const SYMBOL_LINE_BYTES_MAX: usize = 240;
pub const TOKEN_ESTIMATE_MAX: u64 = 512;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum FallbackReason {
    UnknownDialect,
    MalformedHeader,
    MissingRequiredField,
    InvalidRequiredField,
    TokenEstimateTooLarge,
    HeaderTooLarge,
    SymbolLineTooLarge,
    FrameTooLarge,
    InvalidEscape,
    InvalidControlCharacter,
    UnknownSymbol,
    InvalidCacheReuse,
    InvalidDelta,
    UnsupportedDelta,
    MissingCacheEntry,
    DeltaPayloadMismatch,
    UnsafePayload,
}

impl FallbackReason {
    pub fn as_str(self) -> &'static str {
        match self {
            Self::UnknownDialect => "unknown_dialect",
            Self::MalformedHeader => "malformed_header",
            Self::MissingRequiredField => "missing_required_field",
            Self::InvalidRequiredField => "invalid_required_field",
            Self::TokenEstimateTooLarge => "token_estimate_too_large",
            Self::HeaderTooLarge => "header_too_large",
            Self::SymbolLineTooLarge => "symbol_line_too_large",
            Self::FrameTooLarge => "frame_too_large",
            Self::InvalidEscape => "invalid_escape",
            Self::InvalidControlCharacter => "invalid_control_character",
            Self::UnknownSymbol => "unknown_symbol",
            Self::InvalidCacheReuse => "invalid_cache_reuse",
            Self::InvalidDelta => "invalid_delta",
            Self::UnsupportedDelta => "unsupported_delta",
            Self::MissingCacheEntry => "missing_cache_entry",
            Self::DeltaPayloadMismatch => "delta_payload_mismatch",
            Self::UnsafePayload => "unsafe_payload",
        }
    }
}

impl fmt::Display for FallbackReason {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str(self.as_str())
    }
}

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

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Action {
    Continue,
    Refresh,
    WrapUp,
    Stop,
    Escalate,
}

impl Action {
    pub fn from_source_action(action: &str) -> Result<Self, FallbackReason> {
        match action {
            "continue" => Ok(Self::Continue),
            "refresh" => Ok(Self::Refresh),
            "wrap_up" => Ok(Self::WrapUp),
            "stop" => Ok(Self::Stop),
            "escalate" => Ok(Self::Escalate),
            _ => Err(FallbackReason::InvalidRequiredField),
        }
    }

    fn from_code(code: &str) -> Result<Self, FallbackReason> {
        match code {
            "C" => Ok(Self::Continue),
            "R" => Ok(Self::Refresh),
            "W" => Ok(Self::WrapUp),
            "S" => Ok(Self::Stop),
            "E" => Ok(Self::Escalate),
            _ => Err(FallbackReason::InvalidRequiredField),
        }
    }

    pub fn code(self) -> &'static str {
        match self {
            Self::Continue => "C",
            Self::Refresh => "R",
            Self::WrapUp => "W",
            Self::Stop => "S",
            Self::Escalate => "E",
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Mode {
    Startup,
    Steady,
    Delta,
}

impl Mode {
    fn from_str(mode: &str) -> Result<Self, FallbackReason> {
        match mode {
            "startup" => Ok(Self::Startup),
            "steady" => Ok(Self::Steady),
            "delta" => Ok(Self::Delta),
            _ => Err(FallbackReason::InvalidRequiredField),
        }
    }

    pub fn as_str(self) -> &'static str {
        match self {
            Self::Startup => "startup",
            Self::Steady => "steady",
            Self::Delta => "delta",
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Symbol {
    Task,
    Next,
    Rule,
    Memory,
    Blocker,
    Question,
}

impl Symbol {
    pub fn from_char(symbol: char) -> Result<Self, FallbackReason> {
        match symbol {
            'T' => Ok(Self::Task),
            'N' => Ok(Self::Next),
            'R' => Ok(Self::Rule),
            'M' => Ok(Self::Memory),
            '!' => Ok(Self::Blocker),
            '?' => Ok(Self::Question),
            _ => Err(FallbackReason::UnknownSymbol),
        }
    }

    pub fn as_char(self) -> char {
        match self {
            Self::Task => 'T',
            Self::Next => 'N',
            Self::Rule => 'R',
            Self::Memory => 'M',
            Self::Blocker => '!',
            Self::Question => '?',
        }
    }

    fn order(self) -> usize {
        match self {
            Self::Task => 0,
            Self::Blocker => 1,
            Self::Question => 2,
            Self::Next => 3,
            Self::Rule => 4,
            Self::Memory => 5,
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum DeltaMarker {
    Add,
    Repeat,
}

impl DeltaMarker {
    fn from_char(marker: char) -> Result<Self, FallbackReason> {
        match marker {
            '+' => Ok(Self::Add),
            '=' => Ok(Self::Repeat),
            _ => Err(FallbackReason::InvalidDelta),
        }
    }

    fn as_char(self) -> char {
        match self {
            Self::Add => '+',
            Self::Repeat => '=',
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SymbolLine {
    pub marker: Option<DeltaMarker>,
    pub symbol: Symbol,
    pub payload: String,
}

impl SymbolLine {
    pub fn new(symbol: Symbol, payload: String) -> Self {
        Self {
            marker: None,
            symbol,
            payload,
        }
    }

    pub fn with_marker(marker: DeltaMarker, symbol: Symbol, payload: String) -> Self {
        Self {
            marker: Some(marker),
            symbol,
            payload,
        }
    }

    fn render(&self) -> Result<String, FallbackReason> {
        let mut line = String::new();
        if let Some(marker) = self.marker {
            line.push(marker.as_char());
        }
        line.push(self.symbol.as_char());
        line.push(' ');
        line.push_str(&self.payload);
        if line.len() > SYMBOL_LINE_BYTES_MAX {
            return Err(FallbackReason::SymbolLineTooLarge);
        }
        Ok(line)
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Frame {
    pub source_fingerprint: String,
    pub action: Action,
    pub token_estimate: u64,
    pub mode: Mode,
    pub same: bool,
    pub symbols: Vec<SymbolLine>,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct FrameMetrics {
    pub bytes: usize,
    pub token_estimate: u64,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RenderedFrame {
    pub frame: String,
    pub metrics: FrameMetrics,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RenderRequest {
    pub source_fingerprint: String,
    pub action: Action,
    pub token_estimate: u64,
    pub mode: Mode,
    pub same: bool,
    pub symbols: Vec<SymbolLine>,
}

impl RenderRequest {
    pub fn from_runtime_value(
        value: &JsonValue,
        token_estimate: u64,
    ) -> Result<Self, FallbackReason> {
        let source_fingerprint = required_string(value, "source_fingerprint")?.to_owned();
        let action = Action::from_source_action(required_string(value, "action")?)?;
        let mode = value
            .get("mode")
            .and_then(JsonValue::as_str)
            .map(Mode::from_str)
            .transpose()?
            .unwrap_or(Mode::Steady);
        let same = value
            .get("same")
            .and_then(JsonValue::as_bool)
            .unwrap_or(false);
        let symbols = if same {
            Vec::new()
        } else if mode == Mode::Delta {
            delta_lines(value)?
        } else {
            steady_lines(value)?
        };

        Ok(Self {
            source_fingerprint,
            action,
            token_estimate,
            mode,
            same,
            symbols,
        })
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CacheOutcome {
    pub reuses_cached_frame: bool,
    pub applied_symbols: Vec<SymbolLine>,
}

pub fn render_frame(request: &RenderRequest) -> Result<RenderedFrame, FallbackReason> {
    if !valid_fingerprint(&request.source_fingerprint) {
        return Err(FallbackReason::InvalidRequiredField);
    }
    if request.token_estimate > TOKEN_ESTIMATE_MAX {
        return Err(FallbackReason::TokenEstimateTooLarge);
    }
    if request.same && (!request.symbols.is_empty() || request.mode == Mode::Delta) {
        return Err(FallbackReason::InvalidCacheReuse);
    }
    if request.mode == Mode::Delta && request.symbols.iter().any(|line| line.marker.is_none()) {
        return Err(FallbackReason::InvalidDelta);
    }
    if request.mode != Mode::Delta && request.symbols.iter().any(|line| line.marker.is_some()) {
        return Err(FallbackReason::InvalidDelta);
    }

    let mut header = format!(
        "{DIALECT} fp={} act={} tok={} mode={}",
        request.source_fingerprint,
        request.action.code(),
        request.token_estimate,
        request.mode.as_str()
    );
    if request.same {
        header.push_str(" same=1");
    }
    if header.len() > HEADER_BYTES_MAX {
        return Err(FallbackReason::HeaderTooLarge);
    }

    let mut frame = header;
    frame.push('\n');
    for line in &request.symbols {
        frame.push_str(&line.render()?);
        frame.push('\n');
    }
    if frame.len() > FRAME_BYTES_MAX {
        return Err(FallbackReason::FrameTooLarge);
    }

    Ok(RenderedFrame {
        metrics: FrameMetrics {
            bytes: frame.len(),
            token_estimate: request.token_estimate,
        },
        frame,
    })
}

pub fn parse_frame(frame: &str) -> Result<Frame, FallbackReason> {
    if frame.len() > FRAME_BYTES_MAX {
        return Err(FallbackReason::FrameTooLarge);
    }
    let frame = frame.strip_suffix('\n').unwrap_or(frame);
    let mut lines = frame.split('\n');
    let header = lines.next().unwrap_or("");
    if header.len() > HEADER_BYTES_MAX {
        return Err(FallbackReason::HeaderTooLarge);
    }

    let mut fields = header.split_ascii_whitespace();
    let dialect = fields.next().ok_or(FallbackReason::MalformedHeader)?;
    if dialect != DIALECT {
        return Err(FallbackReason::UnknownDialect);
    }

    let header_fields: Vec<(&str, &str)> = fields
        .map(|field| field.split_once('=').ok_or(FallbackReason::MalformedHeader))
        .collect::<Result<_, _>>()?;
    let mut values = HashMap::new();
    for (key, value) in &header_fields {
        if key.is_empty() || value.is_empty() || key.contains(char::is_whitespace) {
            return Err(FallbackReason::MalformedHeader);
        }
        values.insert(*key, *value);
    }

    if header_fields.len() < 3
        || header_fields[0].0 != "fp"
        || header_fields[1].0 != "act"
        || header_fields[2].0 != "tok"
    {
        return Err(FallbackReason::MissingRequiredField);
    }

    let source_fingerprint = values
        .get("fp")
        .copied()
        .ok_or(FallbackReason::MissingRequiredField)?;
    if !valid_fingerprint(source_fingerprint) {
        return Err(FallbackReason::InvalidRequiredField);
    }

    let action = Action::from_code(
        values
            .get("act")
            .copied()
            .ok_or(FallbackReason::MissingRequiredField)?,
    )?;
    let token_estimate = values
        .get("tok")
        .copied()
        .ok_or(FallbackReason::MissingRequiredField)?
        .parse::<u64>()
        .map_err(|_| FallbackReason::InvalidRequiredField)?;
    if token_estimate > TOKEN_ESTIMATE_MAX {
        return Err(FallbackReason::TokenEstimateTooLarge);
    }

    let mode = Mode::from_str(values.get("mode").copied().unwrap_or("steady"))?;
    let same = match values.get("same").copied().unwrap_or("0") {
        "0" => false,
        "1" => true,
        _ => return Err(FallbackReason::InvalidRequiredField),
    };
    if same && mode == Mode::Delta {
        return Err(FallbackReason::InvalidCacheReuse);
    }

    let symbols = lines
        .filter(|line| !line.is_empty())
        .map(|line| parse_symbol_line(line, mode))
        .collect::<Result<Vec<_>, _>>()?;
    if same && !symbols.is_empty() {
        return Err(FallbackReason::InvalidCacheReuse);
    }

    Ok(Frame {
        source_fingerprint: source_fingerprint.to_owned(),
        action,
        token_estimate,
        mode,
        same,
        symbols,
    })
}

pub fn apply_cached_frame(
    incoming: &str,
    cached: Option<&str>,
    delta_supported: bool,
) -> Result<CacheOutcome, FallbackReason> {
    let incoming = parse_frame(incoming)?;
    if incoming.mode == Mode::Delta && !delta_supported {
        return Err(FallbackReason::UnsupportedDelta);
    }
    if incoming.same {
        let cached = cached.ok_or(FallbackReason::MissingCacheEntry)?;
        let cached = parse_frame(cached)?;
        if cached.source_fingerprint != incoming.source_fingerprint {
            return Err(FallbackReason::MissingCacheEntry);
        }
        return Ok(CacheOutcome {
            reuses_cached_frame: true,
            applied_symbols: cached.symbols,
        });
    }
    if incoming.mode == Mode::Delta {
        let cached = cached.ok_or(FallbackReason::MissingCacheEntry)?;
        let cached = parse_frame(cached)?;
        if cached.source_fingerprint != incoming.source_fingerprint {
            return Err(FallbackReason::MissingCacheEntry);
        }
        return Ok(CacheOutcome {
            reuses_cached_frame: false,
            applied_symbols: apply_delta_symbols(&cached.symbols, &incoming.symbols)?,
        });
    }
    Ok(CacheOutcome {
        reuses_cached_frame: false,
        applied_symbols: incoming.symbols,
    })
}

pub fn validate_safe_payload(payload: &str) -> Result<(), FallbackReason> {
    let lower = payload.to_ascii_lowercase();
    let unsafe_payload = lower.contains("ghp_")
        || lower.contains("private key")
        || lower.contains("cookie")
        || lower.contains("without approval")
        || lower.contains("hidden scratchpad")
        || (lower.contains("user:") && lower.contains("assistant:"));
    if unsafe_payload {
        Err(FallbackReason::UnsafePayload)
    } else {
        Ok(())
    }
}

fn valid_fingerprint(fp: &str) -> bool {
    (7..=64).contains(&fp.len())
        && fp
            .bytes()
            .all(|byte| byte.is_ascii_hexdigit() && !byte.is_ascii_uppercase())
}

fn parse_symbol_line(line: &str, mode: Mode) -> Result<SymbolLine, FallbackReason> {
    if line.len() > SYMBOL_LINE_BYTES_MAX {
        return Err(FallbackReason::SymbolLineTooLarge);
    }
    let mut chars = line.chars();
    let first = chars.next().ok_or(FallbackReason::UnknownSymbol)?;
    let (marker, symbol) = if first == '+' || first == '=' {
        let symbol = chars.next().ok_or(FallbackReason::UnknownSymbol)?;
        (Some(DeltaMarker::from_char(first)?), symbol)
    } else {
        (None, first)
    };
    let symbol = Symbol::from_char(symbol)?;

    if mode == Mode::Delta && marker.is_none() {
        return Err(FallbackReason::InvalidDelta);
    }
    if mode != Mode::Delta && marker.is_some() {
        return Err(FallbackReason::InvalidDelta);
    }

    let rest = chars.as_str();
    if !rest.starts_with(' ') {
        return Err(FallbackReason::MalformedHeader);
    }
    let payload = &rest[1..];
    validate_escapes(payload)?;

    Ok(SymbolLine {
        marker,
        symbol,
        payload: payload.to_owned(),
    })
}

fn validate_escapes(payload: &str) -> Result<(), FallbackReason> {
    let mut chars = payload.chars();
    while let Some(ch) = chars.next() {
        if ch != '\\' {
            if ch.is_control() {
                return Err(FallbackReason::InvalidControlCharacter);
            }
            continue;
        }
        match chars.next() {
            Some('\\' | 'n' | 'r' | 't') => {}
            Some('x') => {
                let hi = chars.next().ok_or(FallbackReason::InvalidEscape)?;
                let lo = chars.next().ok_or(FallbackReason::InvalidEscape)?;
                if !hi.is_ascii_hexdigit() || !lo.is_ascii_hexdigit() {
                    return Err(FallbackReason::InvalidEscape);
                }
            }
            _ => return Err(FallbackReason::InvalidEscape),
        }
    }
    Ok(())
}

fn apply_delta_symbols(
    cached: &[SymbolLine],
    delta: &[SymbolLine],
) -> Result<Vec<SymbolLine>, FallbackReason> {
    for line in delta {
        if line.marker == Some(DeltaMarker::Repeat)
            && !cached.iter().any(|cached_line| {
                cached_line.symbol == line.symbol && cached_line.payload == line.payload
            })
        {
            return Err(FallbackReason::DeltaPayloadMismatch);
        }
    }

    let replaced_families: Vec<Symbol> = delta
        .iter()
        .filter(|line| line.marker == Some(DeltaMarker::Add))
        .map(|line| line.symbol)
        .collect();
    let mut merged: Vec<SymbolLine> = cached
        .iter()
        .filter(|line| !replaced_families.contains(&line.symbol))
        .cloned()
        .collect();
    merged.extend(
        delta
            .iter()
            .filter(|line| line.marker == Some(DeltaMarker::Add))
            .map(|line| SymbolLine {
                marker: None,
                symbol: line.symbol,
                payload: line.payload.clone(),
            }),
    );
    merged.sort_by_key(|line| line.symbol.order());
    Ok(merged)
}

fn steady_lines(value: &JsonValue) -> Result<Vec<SymbolLine>, FallbackReason> {
    let mut lines = Vec::new();
    if let Some(title) = optional_string(value, "title")? {
        push_line(&mut lines, Symbol::Task, title)?;
    }
    for payload in string_array(value, "blockers")? {
        push_line(&mut lines, Symbol::Blocker, payload)?;
    }
    for payload in string_array(value, "questions")? {
        push_line(&mut lines, Symbol::Question, payload)?;
    }
    for payload in string_array(value, "next")? {
        push_line(&mut lines, Symbol::Next, payload)?;
    }
    for payload in string_array(value, "rules")? {
        push_line(&mut lines, Symbol::Rule, payload)?;
    }
    for payload in string_array(value, "memory")? {
        push_line(&mut lines, Symbol::Memory, payload)?;
    }
    Ok(lines)
}

fn delta_lines(value: &JsonValue) -> Result<Vec<SymbolLine>, FallbackReason> {
    let delta = value
        .get("delta")
        .and_then(JsonValue::as_array)
        .ok_or(FallbackReason::InvalidDelta)?;
    delta
        .iter()
        .map(|entry| {
            let marker = required_string(entry, "marker")?;
            let marker = marker
                .chars()
                .next()
                .filter(|_| marker.chars().count() == 1)
                .ok_or(FallbackReason::InvalidDelta)
                .and_then(DeltaMarker::from_char)?;
            let symbol = required_string(entry, "symbol")?;
            let symbol = symbol
                .chars()
                .next()
                .filter(|_| symbol.chars().count() == 1)
                .ok_or(FallbackReason::UnknownSymbol)
                .and_then(Symbol::from_char)?;
            let payload = escaped_payload(required_string(entry, "payload")?)?;
            Ok(SymbolLine::with_marker(marker, symbol, payload))
        })
        .collect()
}

fn push_line(
    lines: &mut Vec<SymbolLine>,
    symbol: Symbol,
    payload: &str,
) -> Result<(), FallbackReason> {
    let payload = escaped_payload(payload)?;
    if !payload.is_empty() {
        lines.push(SymbolLine::new(symbol, payload));
    }
    Ok(())
}

fn escaped_payload(payload: &str) -> Result<String, FallbackReason> {
    let normalized = normalize_payload(payload);
    validate_safe_payload(&normalized)?;
    let escaped = escape_payload(&normalized);
    validate_escapes(&escaped)?;
    Ok(escaped)
}

fn normalize_payload(payload: &str) -> String {
    let mut normalized = String::new();
    let mut in_whitespace = false;
    for ch in payload.trim().chars() {
        if ch.is_ascii_whitespace() {
            in_whitespace = true;
            continue;
        }
        if in_whitespace && !normalized.is_empty() {
            normalized.push(' ');
        }
        normalized.push(ch);
        in_whitespace = false;
    }
    normalized
}

fn escape_payload(payload: &str) -> String {
    let mut escaped = String::new();
    for ch in payload.chars() {
        match ch {
            '\\' => escaped.push_str("\\\\"),
            '\n' => escaped.push_str("\\n"),
            '\r' => escaped.push_str("\\r"),
            '\t' => escaped.push_str("\\t"),
            ch if ch.is_control() => escaped.push_str(&format!("\\x{:02X}", ch as u32)),
            ch => escaped.push(ch),
        }
    }
    escaped
}

fn required_string<'a>(value: &'a JsonValue, key: &str) -> Result<&'a str, FallbackReason> {
    value
        .get(key)
        .and_then(JsonValue::as_str)
        .ok_or(FallbackReason::MissingRequiredField)
}

fn optional_string<'a>(value: &'a JsonValue, key: &str) -> Result<Option<&'a str>, FallbackReason> {
    match value.get(key) {
        Some(value) => value
            .as_str()
            .map(Some)
            .ok_or(FallbackReason::InvalidRequiredField),
        None => Ok(None),
    }
}

fn string_array<'a>(value: &'a JsonValue, key: &str) -> Result<Vec<&'a str>, FallbackReason> {
    match value.get(key) {
        Some(JsonValue::Array(values)) => values
            .iter()
            .map(|value| value.as_str().ok_or(FallbackReason::InvalidRequiredField))
            .collect(),
        Some(_) => Err(FallbackReason::InvalidRequiredField),
        None => Ok(Vec::new()),
    }
}