//! Spec extraction. Running the reference function on sampled inputs turns it
//! into a behavioural spec: I/O pairs for synthesis, prefix/hole/suffix for
//! infilling (DATA.md). Inputs are split into shown `examples` and held-out
//! `tests`; correctness is judged on the held-out set.

use rand::Rng;
use rand_chacha::ChaCha20Rng;
use serde::{Deserialize, Serialize};

use crate::grammar::{Program, ProgramFeatures, Type};
use crate::sandbox::{RunError, Sandbox};
use crate::value::LValue;

pub fn type_name(t: Type) -> &'static str {
    match t {
        Type::Number => "number",
        Type::Bool => "bool",
        Type::List => "list",
        Type::Str => "string",
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct IoPair {
    pub input: Vec<LValue>,
    pub output: LValue,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Infill {
    pub prefix: String,
    pub hole: String,
    pub suffix: String,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Features {
    pub lines: usize,
    pub chars: usize,
    pub recursion: bool,
    pub closure: bool,
    pub loops: bool,
    pub table_build: bool,
    pub output_type: String,
}

impl Features {
    pub fn build(source: &str, f: &ProgramFeatures, output: &LValue) -> Features {
        Features {
            lines: source.lines().count(),
            chars: source.len(),
            recursion: f.recursion,
            closure: f.closure,
            loops: f.loops,
            table_build: f.table_build,
            output_type: value_kind(output).to_string(),
        }
    }
}

pub fn value_kind(v: &LValue) -> &'static str {
    match v {
        LValue::Nil => "nil",
        LValue::Bool(_) => "bool",
        LValue::Int(_) => "int",
        LValue::Num(_) => "number",
        LValue::Str(_) => "string",
        LValue::List(_) => "list",
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Record {
    pub id: String,
    pub difficulty: u8,
    pub seed: u64,
    pub source: String,
    pub params: Vec<String>,
    pub examples: Vec<IoPair>, // shown spec (synthesis)
    pub tests: Vec<IoPair>,    // held-out, used for verification
    pub infill: Infill,        // infilling spec
    pub features: Features,
}

#[derive(Debug, Default, Clone)]
pub struct RejectStats {
    pub load_fail: usize,
    pub run_fail: usize,
    pub nonobservable: usize,
    pub nondeterministic: usize,
    pub noop: usize,
    pub duplicate: usize,
    pub weak_tests: usize,
}

/// Build a spec from a program, or return why it was rejected.
pub fn extract(
    sb: &Sandbox,
    rng: &mut ChaCha20Rng,
    prog: &Program,
    n_examples: usize,
    n_tests: usize,
    budget: i64,
    stats: &mut RejectStats,
) -> Option<(Vec<IoPair>, Vec<IoPair>)> {
    if let Err(e) = sb.load_program(&prog.source, budget) {
        stats.load_fail += 1;
        if std::env::var("ECHO_DEBUG").is_ok() {
            eprintln!("LOADFAIL {e:?}\n{}\n---", prog.source);
        }
        return None;
    }

    let total = n_examples + n_tests;
    let mut pairs: Vec<IoPair> = Vec::with_capacity(total);
    let mut seen_inputs: Vec<Vec<LValue>> = Vec::new();
    let mut tries = 0usize;
    let max_tries = total * 8 + 16;

    while pairs.len() < total && tries < max_tries {
        tries += 1;
        let input: Vec<LValue> = prog.params.iter().map(|t| sample(*t, rng)).collect();
        if seen_inputs.contains(&input) {
            continue; // want distinct inputs
        }

        let out = match sb.call_f(&input, budget) {
            Ok(v) => v,
            Err(RunError::Budget) => {
                stats.run_fail += 1;
                return None; // nontermination on a valid input kills the program
            }
            Err(RunError::Lua(_)) => {
                stats.run_fail += 1;
                return None;
            }
            Err(RunError::NonObservable) => {
                stats.nonobservable += 1;
                return None;
            }
        };

        // Determinism check: same input must give the same output.
        match sb.call_f(&input, budget) {
            Ok(v2) if v2 == out => {}
            Ok(_) => {
                stats.nondeterministic += 1;
                return None;
            }
            Err(_) => {
                stats.run_fail += 1;
                return None;
            }
        }

        seen_inputs.push(input.clone());
        pairs.push(IoPair { input, output: out });
    }

    if pairs.len() < total {
        // Couldn't find enough distinct inputs (e.g. tiny domain). Drop quietly.
        stats.noop += 1;
        return None;
    }

    // No-op / constant filter: require at least two distinct outputs.
    let first = &pairs[0].output;
    if pairs.iter().all(|p| &p.output == first) {
        stats.noop += 1;
        return None;
    }

    let tests = pairs.split_off(n_examples);
    Some((pairs, tests))
}

/// Sample `n` input tuples for a given parameter signature (used for mutation
/// probes).
pub fn sample_inputs(params: &[Type], n: usize, rng: &mut ChaCha20Rng) -> Vec<Vec<LValue>> {
    (0..n)
        .map(|_| params.iter().map(|t| sample(*t, rng)).collect())
        .collect()
}

fn sample(t: Type, rng: &mut ChaCha20Rng) -> LValue {
    match t {
        Type::Number => LValue::Int(rng.gen_range(-10..=10)),
        Type::Bool => LValue::Bool(rng.gen_bool(0.5)),
        Type::List => {
            let len = rng.gen_range(1..=6);
            LValue::List((0..len).map(|_| LValue::Int(rng.gen_range(-9..=9))).collect())
        }
        Type::Str => {
            const ALPHA: &[u8] = b"abcdefghijklmnopqrstuvwxyz";
            let len = rng.gen_range(1..=6);
            let s: String = (0..len)
                .map(|_| ALPHA[rng.gen_range(0..ALPHA.len())] as char)
                .collect();
            LValue::Str(s)
        }
    }
}

/// Split source into prefix/hole/suffix by masking one body line (infilling).
pub fn make_infill(source: &str, rng: &mut ChaCha20Rng) -> Infill {
    let lines: Vec<&str> = source.lines().collect();
    // Body lines are everything except the `function f(..)` header and final `end`.
    let candidates: Vec<usize> = (0..lines.len())
        .filter(|&i| {
            let t = lines[i].trim();
            !t.is_empty()
                && !t.starts_with("function f(")
                && t != "end"
                && !t.starts_with("return") // keep a non-trivial structural hole
        })
        .collect();
    let pick = if candidates.is_empty() {
        lines.len().saturating_sub(2) // fall back to the return line
    } else {
        candidates[rng.gen_range(0..candidates.len())]
    };

    let join = |slice: &[&str]| slice.join("\n");
    let prefix = join(&lines[..pick]);
    let hole = lines[pick].to_string();
    let suffix = join(&lines[pick + 1..]);
    Infill { prefix, hole, suffix }
}