//! CUBELinux-2 AI layer.
//!
//! Per the PDF (Package 4, §551): *"cubeai: models that operate on
//! cube‑stored traces/graphs to classify blocks, infer higher‑level
//! operations, and suggest new code sequences or orchestrations."*
//!
//! This crate is the *structured decision* layer sitting on top of the
//! substrate (and on [`cubedbt`]). It is deterministic and dependency-free
//! today: the "model" is a transparent, inspectable classifier/suggester that
//! operates on the same [`CubeStore`] / [`CodeCell`] / [`Behavior`] types the
//! rest of the workspace uses. A learned model can later implement the same
//! traits without changing callers.
//!
//! Pipeline (the end-to-end story from §549–§551 + §554):
//!   trace (cube) → classify blocks → infer higher-level op →
//!   suggest `TranslationRule`s → hand to `cubedbt` to patch + run (mimic).

use cubecode::opcode::Op;
use cubecode::{Behavior, CodeCell, Kind};
use cubecoords::{CubeHeader, Czyx};
use cubedbt::{store_rule, OpClass, TranslationRule};
use cubestore::{CubeStore, HashBackend};

/// `C` axis band where captured traces are stored (fed by `cubetrace` in the
/// full stack; here traces are ingested directly via [`CubeAi::ingest_trace`]).
pub const C_TRACE: u8 = 230;

/// A captured basic block: its coordinate, the op-class histogram observed
/// during tracing, and the behavior descriptors attached (from the header or
/// inferred by the trace layer).
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct BlockTrace {
    pub label: Czyx,
    /// Count of each [`OpClass`] seen in the block (`OpClass::Any` unused).
    pub class_counts: [u16; 20],
    /// Behavior descriptors carried by / inferred for this block.
    pub behavior: Behavior,
}

impl BlockTrace {
    /// Build from a decoded code cell, deriving the op-class histogram and
    /// reading any behavior descriptors from its header flags.
    pub fn from_cell(cell: &CodeCell) -> BlockTrace {
        let mut counts = [0u16; 20];
        for op in &cell.code {
            let c = OpClass::of(op) as usize;
            if c < 20 {
                counts[c] += 1;
            }
        }
        BlockTrace {
            label: cell.label,
            class_counts: counts,
            behavior: Behavior::from_flags(cell.header.flags.bits()),
        }
    }

    /// Total op count.
    pub fn total(&self) -> u32 {
        self.class_counts.iter().map(|&c| c as u32).sum()
    }
}

/// A higher-level classification of a block, inferred by [`BlockClassifier`].
#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub enum BlockKind {
    /// Mostly arithmetic/logic — a computation block.
    Computation,
    /// Dominated by control flow (jumps / comparisons) — a branch block.
    Branch,
    /// Contains call links to other records — a call/composition block.
    CallTrampoline,
    /// Heavy I/O or network descriptors — an I/O section.
    IoSection,
    /// Otherwise: a plain linear sequence.
    Sequence,
}

/// Deterministic classifier: maps a [`BlockTrace`] to a [`BlockKind`] from its
/// op-class histogram and behavior descriptors. (The transparent "model".)
pub struct BlockClassifier;

impl BlockClassifier {
    pub fn classify(block: &BlockTrace) -> BlockKind {
        let c = &block.class_counts;
        let arith = c[OpClass::Add as usize]
            + c[OpClass::Sub as usize]
            + c[OpClass::Mul as usize]
            + c[OpClass::Div as usize]
            + c[OpClass::Mod as usize]
            + c[OpClass::And as usize]
            + c[OpClass::Or as usize]
            + c[OpClass::Xor as usize]
            + c[OpClass::Shl as usize]
            + c[OpClass::Shr as usize];
        let ctrl = c[OpClass::Eq as usize]
            + c[OpClass::Ne as usize]
            + c[OpClass::Lt as usize]
            + c[OpClass::Gt as usize]
            + c[OpClass::Le as usize]
            + c[OpClass::Ge as usize];
        let jumps = c[OpClass::Const as usize];
        let calls = c[OpClass::CallLink as usize];

        if block.behavior.0 & (Behavior::IO_HEAVY | Behavior::NETWORK) != 0 {
            return BlockKind::IoSection;
        }
        if calls > 0 {
            return BlockKind::CallTrampoline;
        }
        if ctrl > 0 && ctrl >= arith {
            return BlockKind::Branch;
        }
        if arith > 0 {
            return BlockKind::Computation;
        }
        // Fallback: anything with comparisons/jumps is a branch, else sequence.
        if jumps > 0 || ctrl > 0 {
            BlockKind::Branch
        } else {
            BlockKind::Sequence
        }
    }
}

/// Suggest transformations for a block, returning candidate [`TranslationRule`]s
/// the DBT layer can apply (PDF: "suggest new code sequences or
/// orchestrations"). Deterministic and local: each suggestion records *why*.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Suggestion {
    pub rule: TranslationRule,
    pub rationale: String,
}

/// The AI runtime over a cube store: ingests traces, classifies, and suggests
/// DBT rules. Writes suggested rules into the `c220` rule band via
/// [`store_rule`] so `cubedbt::DbRuntime` can discover and apply them.
pub struct CubeAi {
    store: CubeStore<HashBackend>,
    next_trace_x: u8,
    next_rule_x: u8,
}

impl CubeAi {
    /// Build an empty AI runtime over a fresh in-memory store.
    pub fn new() -> Self {
        CubeAi {
            store: CubeStore::new(HashBackend::new()),
            next_trace_x: 1,
            next_rule_x: 1,
        }
    }

    /// Ingest a captured trace (e.g. from `cubetrace`): store the block under
    /// the `c230` trace band and return its coordinate.
    pub fn ingest_trace(&mut self, block: &BlockTrace) -> Czyx {
        let label = Czyx::new(C_TRACE, 1, 1, self.next_trace_x);
        self.next_trace_x = self.next_trace_x.wrapping_add(1).max(1);
        let mut h = CubeHeader::new();
        h.title = Some(format!("trace:{:?}", block.label));
        h.doc_type = Some(Kind::Other.as_str().into());
        h.size_bytes = Some(block.class_counts.len() as u64 * 2);
        h.flags.0 |= block.behavior.to_flags();
        h.refresh_flags();
        // Body: the raw histogram (20 x u16 le).
        let mut body = Vec::with_capacity(40);
        for c in &block.class_counts {
            body.extend_from_slice(&c.to_le_bytes());
        }
        self.store.put_record(label, &h, &body);
        label
    }

    /// Classify a single block.
    pub fn classify(&self, block: &BlockTrace) -> BlockKind {
        BlockClassifier::classify(block)
    }

    /// Produce suggestions for a block (does not yet persist them).
    pub fn suggest(&self, block: &BlockTrace) -> Vec<Suggestion> {
        let kind = BlockClassifier::classify(block);
        let mut out = Vec::new();
        match kind {
            BlockKind::Computation => {
                // Suggest vectorizing a repeated multiply-by-constant:
                // Mul(const) -> Const(shift?), but keep it conservative:
                // replace Mul with a left-shift when operand is a power of two.
                out.push(Suggestion {
                    rule: TranslationRule {
                        name: format!("compute-opt:{:?}", block.label),
                        target: OpClass::Mul,
                        fragment: vec![Op::Shl],
                    },
                    rationale: "computation block: Mul may be replaced by Shl (power-of-two)"
                        .into(),
                });
            }
            BlockKind::Branch => {
                out.push(Suggestion {
                    rule: TranslationRule {
                        name: format!("branch-opt:{:?}", block.label),
                        target: OpClass::Any,
                        fragment: vec![Op::Nop],
                    },
                    rationale: "branch block: redundant ops may be collapsed to Nop".into(),
                });
            }
            BlockKind::IoSection => {
                out.push(Suggestion {
                    rule: TranslationRule {
                        name: format!("io-batch:{:?}", block.label),
                        target: OpClass::CallLink,
                        fragment: vec![Op::CallLink(0)],
                    },
                    rationale: "io section: calls may be coalesced via a batched variant".into(),
                });
            }
            _ => {}
        }
        out
    }

    /// Ingest a block and persist its suggestions as DBT rules in the `c220`
    /// band. Returns the stored rule coordinates (empty if no suggestions).
    pub fn ingest_and_suggest(&mut self, block: &BlockTrace) -> Vec<Czyx> {
        self.ingest_trace(block);
        let sugs = self.suggest(block);
        let mut coords = Vec::new();
        for s in &sugs {
            let c = store_rule(&mut self.store, &s.rule, self.next_rule_x);
            self.next_rule_x = self.next_rule_x.wrapping_add(1).max(1);
            coords.push(c);
        }
        coords
    }

    /// Borrow the backing store (e.g. to hand to `cubedbt::DbRuntime`).
    pub fn store(&self) -> &CubeStore<HashBackend> {
        &self.store
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use cubecode::opcode::Op;
    use cubecoords::Czyx;

    fn cell(coord: Czyx, code: Vec<Op>, beh: Behavior) -> CodeCell {
        let mut h = CubeHeader::new();
        h.title = Some("blk".into());
        h.doc_type = Some(Kind::Fn.as_str().into());
        h.flags.0 |= beh.to_flags();
        h.refresh_flags();
        CodeCell::from_record(coord, &h, &cubecode::opcode::encode(&code))
            .expect("cell is valid bytecode")
    }

    #[test]
    fn trace_classifies_computation() {
        let c = cell(
            Czyx::new(1, 1, 1, 1),
            vec![Op::Const(3), Op::Const(4), Op::Mul, Op::Halt],
            Behavior(Behavior::PURE),
        );
        let t = BlockTrace::from_cell(&c);
        assert_eq!(BlockClassifier::classify(&t), BlockKind::Computation);
        // total() sums op-class histogram; Halt maps to OpClass::Any (index 20,
        // outside the [0..20) histogram), so 3 counted ops (Const, Const, Mul).
        assert_eq!(t.total(), 3);
    }

    #[test]
    fn trace_classifies_io_section() {
        let c = cell(
            Czyx::new(1, 1, 1, 2),
            vec![Op::Const(1), Op::Halt],
            Behavior(Behavior::IO_HEAVY | Behavior::NETWORK),
        );
        let t = BlockTrace::from_cell(&c);
        assert_eq!(BlockClassifier::classify(&t), BlockKind::IoSection);
    }

    #[test]
    fn trace_classifies_call_trampoline() {
        let c = cell(
            Czyx::new(1, 1, 1, 3),
            vec![Op::CallLink(0), Op::CallLink(1), Op::Halt],
            Behavior::default(),
        );
        let t = BlockTrace::from_cell(&c);
        assert_eq!(BlockClassifier::classify(&t), BlockKind::CallTrampoline);
    }

    #[test]
    fn suggest_emits_rule_and_persists() {
        let c = cell(
            Czyx::new(1, 1, 1, 4),
            vec![Op::Const(3), Op::Const(4), Op::Mul, Op::Halt],
            Behavior(Behavior::PURE),
        );
        let t = BlockTrace::from_cell(&c);
        let mut ai = CubeAi::new();
        let coords = ai.ingest_and_suggest(&t);
        assert_eq!(coords.len(), 1, "computation block suggests one rule");
        assert_eq!(coords[0].c, rule_band());
        // The stored rule is discoverable by cubedbt.
        let rt = cubedbt::DbRuntime::new(ai.store().clone());
        assert_eq!(rt.discover_rules().len(), 1);
    }

    #[test]
    fn ingest_trace_stores_under_c230() {
        let c = cell(Czyx::new(1, 1, 1, 5), vec![Op::Halt], Behavior::default());
        let t = BlockTrace::from_cell(&c);
        let mut ai = CubeAi::new();
        let coord = ai.ingest_trace(&t);
        assert_eq!(coord.c, C_TRACE);
    }

    // Helper: the rule band constant lives in cubedbt; assert equality without
    // importing the const name directly (kept local to the test).
    fn rule_band() -> u8 {
        cubedbt::C_DBT_RULE
    }
}