//! Compiler for filter expressions.
//!
//! Compiles an AST into bytecode that can be executed by the VM.

use std::collections::HashMap;

use regex::bytes::Regex;
use thiserror::Error;

use crate::opcode::Opcode;
use crate::parser::{Expr, ParseError, ParserConfig};
use crate::vm::CompiledFilter;

/// Compilation error types.
#[derive(Debug, Clone, Error)]
#[allow(missing_docs)]
pub enum CompileError {
    #[error("Parse error: {0}")]
    Parse(#[from] ParseError),

    #[error("Unknown field '{0}'. Available fields: {1}")]
    UnknownField(String, String),

    #[error("Invalid regex pattern '{pattern}': {error}")]
    InvalidRegex { pattern: String, error: String },

    #[error("Too many strings (max 65535)")]
    TooManyStrings,

    #[error("Too many regexes (max 65535)")]
    TooManyRegexes,

    #[error("Too many string sets (max 65535)")]
    TooManySets,
}

/// Compiler state during bytecode generation.
struct Compiler<'a> {
    config: &'a ParserConfig,
    bytecode: Vec<u8>,
    strings: Vec<Vec<u8>>,
    string_map: HashMap<Vec<u8>, u16>,
    regexes: Vec<Regex>,
    regex_map: HashMap<String, u16>,
    string_sets: Vec<Vec<u16>>,
}

impl<'a> Compiler<'a> {
    fn new(config: &'a ParserConfig) -> Self {
        Self {
            config,
            bytecode: Vec::new(),
            strings: Vec::new(),
            string_map: HashMap::new(),
            regexes: Vec::new(),
            regex_map: HashMap::new(),
            string_sets: Vec::new(),
        }
    }

    /// Intern a string and return its index.
    fn intern_string(&mut self, s: &str) -> Result<u16, CompileError> {
        let bytes = s.as_bytes().to_vec();
        if let Some(&idx) = self.string_map.get(&bytes) {
            return Ok(idx);
        }

        let idx = self.strings.len();
        if idx > u16::MAX as usize {
            return Err(CompileError::TooManyStrings);
        }

        self.string_map.insert(bytes.clone(), idx as u16);
        self.strings.push(bytes);
        Ok(idx as u16)
    }

    /// Intern a regex and return its index.
    fn intern_regex(&mut self, pattern: &str) -> Result<u16, CompileError> {
        if let Some(&idx) = self.regex_map.get(pattern) {
            return Ok(idx);
        }

        let regex = Regex::new(pattern).map_err(|e| CompileError::InvalidRegex {
            pattern: pattern.to_string(),
            error: e.to_string(),
        })?;

        let idx = self.regexes.len();
        if idx > u16::MAX as usize {
            return Err(CompileError::TooManyRegexes);
        }

        self.regex_map.insert(pattern.to_string(), idx as u16);
        self.regexes.push(regex);
        Ok(idx as u16)
    }

    /// Add a string set and return its index.
    fn add_string_set(&mut self, values: &[String]) -> Result<u16, CompileError> {
        let indices: Vec<u16> = values
            .iter()
            .map(|v| self.intern_string(v))
            .collect::<Result<_, _>>()?;

        let idx = self.string_sets.len();
        if idx > u16::MAX as usize {
            return Err(CompileError::TooManySets);
        }

        self.string_sets.push(indices);
        Ok(idx as u16)
    }

    /// Look up a field name and return its part index.
    fn lookup_field(&self, name: &str) -> Result<u8, CompileError> {
        // Try case-insensitive lookup
        let upper = name.to_uppercase();
        self.config
            .fields
            .get(&upper)
            .or_else(|| self.config.fields.get(name))
            .copied()
            .ok_or_else(|| {
                let available: Vec<_> = self.config.fields.keys().cloned().collect();
                CompileError::UnknownField(name.to_string(), available.join(", "))
            })
    }

    /// Emit a single byte.
    fn emit(&mut self, byte: u8) {
        self.bytecode.push(byte);
    }

    /// Emit a u16 in little-endian format.
    fn emit_u16(&mut self, value: u16) {
        self.bytecode.extend_from_slice(&value.to_le_bytes());
    }

    /// Emit an i16 in little-endian format.
    fn emit_i16(&mut self, value: i16) {
        self.bytecode.extend_from_slice(&value.to_le_bytes());
    }

    /// Current bytecode offset (for backpatching).
    fn offset(&self) -> usize {
        self.bytecode.len()
    }

    /// Backpatch an i16 at the given bytecode position.
    fn patch_i16(&mut self, pos: usize, value: i16) {
        let bytes = value.to_le_bytes();
        self.bytecode[pos] = bytes[0];
        self.bytecode[pos + 1] = bytes[1];
    }

    /// Compile an expression — thin dispatcher to per-variant helpers.
    fn compile_expr(&mut self, expr: &Expr) -> Result<(), CompileError> {
        match expr {
            // Literals
            Expr::Bool(val) => self.compile_bool(*val),
            Expr::Rand(n) => self.compile_rand(*n),

            // Payload-wide operations
            Expr::Contains(s) => self.compile_payload_string(Opcode::Contains, s),
            Expr::StartsWith(s) => self.compile_payload_string(Opcode::StartsWith, s),
            Expr::EndsWith(s) => self.compile_payload_string(Opcode::EndsWith, s),
            Expr::Equals(s) => self.compile_payload_string(Opcode::Equals, s),
            Expr::Matches(pattern) => self.compile_payload_matches(pattern),

            // Part-specific operations
            Expr::PartContains { part, value } => {
                self.compile_part_string(Opcode::PartContains, part, value)
            }
            Expr::PartIContains { part, value } => {
                self.compile_part_string(Opcode::PartIContains, part, value)
            }
            Expr::PartStartsWith { part, value } => {
                self.compile_part_string(Opcode::PartStartsWith, part, value)
            }
            Expr::PartEndsWith { part, value } => {
                self.compile_part_string(Opcode::PartEndsWith, part, value)
            }
            Expr::PartEquals { part, value } => {
                self.compile_part_string(Opcode::PartEquals, part, value)
            }
            Expr::PartIEquals { part, value } => {
                self.compile_part_string(Opcode::PartIEquals, part, value)
            }
            Expr::PartNotEquals { part, value } => self.compile_part_not_equals(part, value),
            Expr::PartMatches { part, pattern } => self.compile_part_matches(part, pattern),
            Expr::PartInSet { part, values } => self.compile_part_in_set(part, values),
            Expr::PartIsEmpty { part } => self.compile_part_flag(Opcode::PartIsEmpty, part),
            Expr::PartNotEmpty { part } => self.compile_part_flag(Opcode::PartNotEmpty, part),

            // Header operations
            Expr::HeaderEquals {
                part,
                header,
                value,
            } => self.compile_header_string(Opcode::HeaderEquals, part, header, value),
            Expr::HeaderIEquals {
                part,
                header,
                value,
            } => self.compile_header_string(Opcode::HeaderIEquals, part, header, value),
            Expr::HeaderContains {
                part,
                header,
                value,
            } => self.compile_header_string(Opcode::HeaderContains, part, header, value),
            Expr::HeaderExists { part, header } => self.compile_header_exists(part, header),

            // Boolean operations — short-circuit with jumps
            Expr::And(left, right) => self.compile_short_circuit(Opcode::JumpIfFalse, left, right),
            Expr::Or(left, right) => self.compile_short_circuit(Opcode::JumpIfTrue, left, right),
            Expr::Not(inner) => self.compile_not(inner),
        }
    }

    // ── Literal helpers ──────────────────────────────────────────────

    fn compile_bool(&mut self, val: bool) -> Result<(), CompileError> {
        self.emit(if val {
            Opcode::PushTrue as u8
        } else {
            Opcode::PushFalse as u8
        });
        Ok(())
    }

    fn compile_rand(&mut self, n: u16) -> Result<(), CompileError> {
        self.emit(Opcode::Rand as u8);
        self.emit_u16(n);
        Ok(())
    }

    // ── Payload-wide helpers ─────────────────────────────────────────

    /// Emit an opcode that takes a single interned string (payload-wide ops).
    fn compile_payload_string(&mut self, op: Opcode, s: &str) -> Result<(), CompileError> {
        let idx = self.intern_string(s)?;
        self.emit(op as u8);
        self.emit_u16(idx);
        Ok(())
    }

    fn compile_payload_matches(&mut self, pattern: &str) -> Result<(), CompileError> {
        let idx = self.intern_regex(pattern)?;
        self.emit(Opcode::Matches as u8);
        self.emit_u16(idx);
        Ok(())
    }

    // ── Part-specific helpers ────────────────────────────────────────

    /// Emit an opcode of the form: `<op> <part_idx> <string_idx>`.
    fn compile_part_string(
        &mut self,
        op: Opcode,
        part: &str,
        value: &str,
    ) -> Result<(), CompileError> {
        let part_idx = self.lookup_field(part)?;
        let str_idx = self.intern_string(value)?;
        self.emit(op as u8);
        self.emit(part_idx);
        self.emit_u16(str_idx);
        Ok(())
    }

    /// Compile `!=` as `PartEquals` followed by `Not`.
    fn compile_part_not_equals(&mut self, part: &str, value: &str) -> Result<(), CompileError> {
        self.compile_part_string(Opcode::PartEquals, part, value)?;
        self.emit(Opcode::Not as u8);
        Ok(())
    }

    fn compile_part_matches(&mut self, part: &str, pattern: &str) -> Result<(), CompileError> {
        let part_idx = self.lookup_field(part)?;
        let regex_idx = self.intern_regex(pattern)?;
        self.emit(Opcode::PartMatches as u8);
        self.emit(part_idx);
        self.emit_u16(regex_idx);
        Ok(())
    }

    fn compile_part_in_set(&mut self, part: &str, values: &[String]) -> Result<(), CompileError> {
        let part_idx = self.lookup_field(part)?;
        let set_idx = self.add_string_set(values)?;
        self.emit(Opcode::PartInSet as u8);
        self.emit(part_idx);
        self.emit_u16(set_idx);
        Ok(())
    }

    /// Emit an opcode that takes only a part index (e.g. `PartIsEmpty`, `PartNotEmpty`).
    fn compile_part_flag(&mut self, op: Opcode, part: &str) -> Result<(), CompileError> {
        let part_idx = self.lookup_field(part)?;
        self.emit(op as u8);
        self.emit(part_idx);
        Ok(())
    }

    // ── Header helpers ───────────────────────────────────────────────

    /// Emit an opcode of the form: `<op> <part_idx> <header_idx> <value_idx>`.
    fn compile_header_string(
        &mut self,
        op: Opcode,
        part: &str,
        header: &str,
        value: &str,
    ) -> Result<(), CompileError> {
        let part_idx = self.lookup_field(part)?;
        let hdr_idx = self.intern_string(header)?;
        let val_idx = self.intern_string(value)?;
        self.emit(op as u8);
        self.emit(part_idx);
        self.emit_u16(hdr_idx);
        self.emit_u16(val_idx);
        Ok(())
    }

    fn compile_header_exists(&mut self, part: &str, header: &str) -> Result<(), CompileError> {
        let part_idx = self.lookup_field(part)?;
        let hdr_idx = self.intern_string(header)?;
        self.emit(Opcode::HeaderExists as u8);
        self.emit(part_idx);
        self.emit_u16(hdr_idx);
        Ok(())
    }

    // ── Boolean / logical helpers ────────────────────────────────────

    /// Compile AND/OR with short-circuit jumps.
    fn compile_short_circuit(
        &mut self,
        jump_op: Opcode,
        left: &Expr,
        right: &Expr,
    ) -> Result<(), CompileError> {
        // Emit left operand
        self.compile_expr(left)?;
        // Conditional jump: skip right operand when short-circuit fires
        let opcode_pos = self.offset();
        self.emit(jump_op as u8);
        let patch_pos = self.offset();
        self.emit_i16(0); // placeholder
        // Emit right operand (its result becomes the AND/OR result)
        self.compile_expr(right)?;
        // Backpatch: offset is relative to opcode position (VM does pc += offset)
        let jump_target = self.offset();
        let offset = (jump_target as isize - opcode_pos as isize) as i16;
        self.patch_i16(patch_pos, offset);
        Ok(())
    }

    fn compile_not(&mut self, inner: &Expr) -> Result<(), CompileError> {
        self.compile_expr(inner)?;
        self.emit(Opcode::Not as u8);
        Ok(())
    }

    /// Finish compilation and return the compiled filter.
    fn finish(mut self, source: String, prefilter_literal: Option<Vec<u8>>) -> CompiledFilter {
        self.emit(Opcode::Return as u8);

        CompiledFilter::new_with_prefilter(
            self.bytecode,
            self.strings,
            self.regexes,
            self.string_sets,
            self.config.delimiter.clone(),
            prefilter_literal,
            source,
        )
    }
}

/// Walk the AST and return a byte sequence that any payload satisfying
/// `expr` is guaranteed to contain, or `None` if no sound substring can
/// be derived.
///
/// The returned literal is used as a batch-level SIMD prefilter: callers
/// can run a single memmem scan over a whole batch of record payloads to
/// rule out batches with no chance of matching. False positives (the
/// literal is found but no record actually matches) fall through to the
/// per-record filter, so soundness is the only correctness requirement.
///
/// Selectivity strategy: when a sound literal can be derived from a
/// `FIELD == "value"` clause and FIELD has a positive part index, we
/// prepend the delimiter (the field cannot be at the start of the
/// payload). For top-level AND, we pick the longer of the two branches'
/// candidates. For OR, we return a candidate only if both branches agree
/// on the exact same bytes.
fn derive_required_substring(expr: &Expr, config: &ParserConfig) -> Option<Vec<u8>> {
    use Expr::*;
    match expr {
        // Always-true / non-deterministic / NOT / regex / empty-checks:
        // no sound required substring.
        Bool(_) | Rand(_) | Matches(_) | Not(_) | PartIsEmpty { .. } | PartNotEmpty { .. } => None,

        // Payload-wide literal ops: the literal must appear somewhere.
        Contains(s) | StartsWith(s) | EndsWith(s) | Equals(s) => {
            non_empty_bytes(s.as_bytes())
        }

        // Part-specific positional ops on case-sensitive literals.
        // Soundness: the value appears as (a substring of) some field.
        // Fields are delimited; for any part with index > 0 the payload
        // contains `delimiter + value`. For part 0 the payload starts
        // with `value` (no leading delim).
        PartContains { value, .. } | PartStartsWith { value, .. } | PartEndsWith { value, .. } => {
            non_empty_bytes(value.as_bytes())
        }
        PartEquals { part, value } => part_eq_literal(part, value, config),

        // Case-insensitive ops cannot be turned into a sound byte literal.
        PartIContains { .. } | PartIEquals { .. } | HeaderIEquals { .. } => None,

        // != cannot rule out anything.
        PartNotEquals { .. } => None,

        // PartMatches is regex on a field. Skip.
        PartMatches { .. } => None,

        // FIELD in {"a","b",...}: a sound literal is the common prefix
        // of all values (must appear in any matching payload). If the
        // prefix is empty, no sound literal.
        PartInSet { values, .. } => longest_common_prefix(values)
            .filter(|p| !p.is_empty())
            .map(|p| p.into_bytes()),

        // Header ops with case-sensitive equality on the value.
        // Soundness: the matching header value appears in the payload.
        HeaderEquals { value, .. } | HeaderContains { value, .. } => {
            non_empty_bytes(value.as_bytes())
        }
        HeaderExists { .. } => None,

        // AND: both branches must hold. Either branch's required
        // substring is sufficient. Pick the longer (more selective).
        And(a, b) => {
            let ra = derive_required_substring(a, config);
            let rb = derive_required_substring(b, config);
            match (ra, rb) {
                (Some(x), Some(y)) => Some(if x.len() >= y.len() { x } else { y }),
                (Some(x), None) | (None, Some(x)) => Some(x),
                (None, None) => None,
            }
        }

        // OR: at least one branch must hold, so a sound required
        // substring must appear in EVERY satisfying payload of EITHER
        // branch. We can only return one if both branches yield the same
        // bytes. (A real intersection would need substring-of-substring
        // analysis; not worth the complexity.)
        Or(a, b) => {
            let ra = derive_required_substring(a, config)?;
            let rb = derive_required_substring(b, config)?;
            if ra == rb {
                Some(ra)
            } else {
                None
            }
        }
    }
}

fn non_empty_bytes(b: &[u8]) -> Option<Vec<u8>> {
    if b.is_empty() {
        None
    } else {
        Some(b.to_vec())
    }
}

/// Build a delimiter-contextualized required substring for `FIELD == value`.
/// For non-zero field indexes we can sound-prepend the delimiter (the
/// field is preceded by one). For index 0 (or unknown fields, conservatively)
/// we return just the value.
fn part_eq_literal(part: &str, value: &str, config: &ParserConfig) -> Option<Vec<u8>> {
    let value_bytes = value.as_bytes();
    if value_bytes.is_empty() {
        return None;
    }
    let idx = config.fields.get(part).copied().unwrap_or(0);
    let mut out = Vec::with_capacity(config.delimiter.len() + value_bytes.len());
    if idx > 0 {
        out.extend_from_slice(&config.delimiter);
    }
    out.extend_from_slice(value_bytes);
    Some(out)
}

fn longest_common_prefix(values: &[String]) -> Option<String> {
    let mut iter = values.iter();
    let first = iter.next()?.as_bytes();
    let mut len = first.len();
    for v in iter {
        let vb = v.as_bytes();
        let common = first
            .iter()
            .zip(vb.iter())
            .take_while(|(a, b)| a == b)
            .count();
        if common < len {
            len = common;
        }
        if len == 0 {
            return None;
        }
    }
    Some(String::from_utf8_lossy(&first[..len]).into_owned())
}

/// Compile a filter expression string into a CompiledFilter.
///
/// # Arguments
/// * `source` - The filter expression string
/// * `config` - Parser configuration with field mappings
///
/// # Returns
/// A `CompiledFilter` ready for evaluation.
///
/// # Example
/// ```
/// use bytecode_filter::{compile, ParserConfig};
/// use bytes::Bytes;
///
/// let mut config = ParserConfig::default();
/// config.add_field("STATUS", 0);
/// config.add_field("CODE", 1);
/// let filter = compile(r#"STATUS == "ok""#, &config).unwrap();
///
/// let record = Bytes::from("ok;;;200");
/// assert!(filter.evaluate(record));
/// ```
///
/// # Errors
/// Returns `CompileError` if parsing or compilation fails.
pub fn compile(source: &str, config: &ParserConfig) -> Result<CompiledFilter, CompileError> {
    let expr = crate::parser::parse(source, config)?;
    compile_expr(&expr, config, source.to_string())
}

/// Compile a pre-parsed AST into a CompiledFilter.
///
/// # Errors
/// Returns `CompileError` if the expression contains invalid operations.
pub fn compile_expr(
    expr: &Expr,
    config: &ParserConfig,
    source: String,
) -> Result<CompiledFilter, CompileError> {
    let mut compiler = Compiler::new(config);
    compiler.compile_expr(expr)?;
    let prefilter_literal = derive_required_substring(expr, config);
    Ok(compiler.finish(source, prefilter_literal))
}

#[cfg(test)]
mod tests {
    use bytes::Bytes;

    use super::*;

    fn test_config() -> ParserConfig {
        let mut config = ParserConfig::default();
        config.add_field("LEVEL", 0);
        config.add_field("CODE", 1);
        config.add_field("METHOD", 2);
        config.add_field("PATH", 3);
        config.add_field("HEADERS", 4);
        config.add_field("BODY", 5);
        config
    }

    fn compile_and_test(input: &str, payload: &str, expected: bool) {
        let config = test_config();
        let filter = compile(input, &config).expect("Failed to compile");
        let result = filter.evaluate(Bytes::from(payload.to_string()));
        assert_eq!(
            result, expected,
            "Filter '{}' on payload '{}' expected {} but got {}",
            input, payload, expected, result
        );
    }

    #[test]
    fn test_compile_true() {
        compile_and_test("true", "", true);
    }

    #[test]
    fn test_compile_false() {
        compile_and_test("false", "", false);
    }

    #[test]
    fn test_compile_payload_contains() {
        compile_and_test(r#"payload contains "error""#, "an error occurred", true);
        compile_and_test(r#"payload contains "error""#, "all good", false);
    }

    #[test]
    fn test_compile_field_equals() {
        compile_and_test(r#"LEVEL == "error""#, "error;;;500;;;GET", true);
        compile_and_test(r#"LEVEL == "error""#, "info;;;500;;;GET", false);
    }

    #[test]
    fn test_compile_field_in_set() {
        compile_and_test(r#"LEVEL in {"error", "warn", "fatal"}"#, "error;;;500;;;GET", true);
        compile_and_test(r#"LEVEL in {"error", "warn", "fatal"}"#, "warn;;;500;;;GET", true);
        compile_and_test(r#"LEVEL in {"error", "warn", "fatal"}"#, "info;;;500;;;GET", false);
    }

    #[test]
    fn test_compile_and() {
        compile_and_test(
            r#"LEVEL == "error" AND CODE == "500""#,
            "error;;;500;;;GET",
            true,
        );
        compile_and_test(
            r#"LEVEL == "error" AND CODE == "500""#,
            "error;;;200;;;GET",
            false,
        );
    }

    #[test]
    fn test_compile_or() {
        compile_and_test(
            r#"LEVEL == "error" OR LEVEL == "warn""#,
            "error;;;500;;;GET",
            true,
        );
        compile_and_test(
            r#"LEVEL == "error" OR LEVEL == "warn""#,
            "warn;;;500;;;GET",
            true,
        );
        compile_and_test(
            r#"LEVEL == "error" OR LEVEL == "warn""#,
            "info;;;500;;;GET",
            false,
        );
    }

    #[test]
    fn test_compile_not() {
        compile_and_test(r#"NOT LEVEL == "debug""#, "error;;;500;;;GET", true);
        compile_and_test(r#"NOT LEVEL == "debug""#, "debug;;;500;;;GET", false);
    }

    #[test]
    fn test_compile_header_iequals() {
        let mut parts = vec![""; 6];
        parts[0] = "error";
        parts[4] = "X-Custom: value\r\n";
        let payload = parts.join(";;;");

        let config = test_config();
        let filter = compile(
            r#"HEADERS.header("x-custom") iequals "value""#,
            &config,
        )
        .unwrap();

        assert!(filter.evaluate(Bytes::from(payload)));
    }

    #[test]
    fn test_compile_complex_filter() {
        let filter_str = r#"
            LEVEL == "error"
            AND CODE == "500"
            AND HEADERS.header("Content-Type") iequals "application/json"
        "#;

        let config = test_config();
        let filter = compile(filter_str, &config).unwrap();

        let mut parts = vec![""; 6];
        parts[0] = "error";
        parts[1] = "500";
        parts[4] = "Content-Type: application/json\r\n";
        let payload = parts.join(";;;");
        assert!(filter.evaluate(Bytes::from(payload)));

        parts[0] = "info";
        let payload = parts.join(";;;");
        assert!(!filter.evaluate(Bytes::from(payload)));
    }

    #[test]
    fn test_compile_rand() {
        crate::vm::reset_rand_counter();

        let config = test_config();
        let filter = compile("rand(2)", &config).unwrap();

        assert!(filter.evaluate(Bytes::new()));
        assert!(!filter.evaluate(Bytes::new()));
        assert!(filter.evaluate(Bytes::new()));
        assert!(!filter.evaluate(Bytes::new()));
    }

    #[test]
    fn test_compile_regex() {
        compile_and_test(r#"payload matches "error_[0-9]+""#, "found error_123", true);
        compile_and_test(r#"payload matches "error_[0-9]+""#, "no errors", false);
    }

    #[test]
    fn test_compile_unknown_field() {
        let config = test_config();
        let result = compile(r#"UNKNOWN_FIELD == "x""#, &config);
        assert!(matches!(result, Err(CompileError::UnknownField(_, _))));
    }

    #[test]
    fn test_compile_invalid_regex() {
        let config = test_config();
        let result = compile(r#"payload matches "[invalid""#, &config);
        assert!(matches!(result, Err(CompileError::InvalidRegex { .. })));
    }

    #[test]
    fn test_bytecode_size() {
        let config = test_config();

        let filter = compile(r#"LEVEL == "error""#, &config).unwrap();
        assert_eq!(filter.bytecode_len(), 5); // PartEquals(1 + 1 + 2) + Return(1)

        let filter = compile(
            r#"LEVEL == "error" AND CODE == "500""#,
            &config,
        )
        .unwrap();
        assert_eq!(filter.bytecode_len(), 12); // 2x PartEquals(4) + JumpIfFalse(3) + Return(1)
    }
}
