use std::collections::HashMap;
use std::str::FromStr;

use crate::ir::{BinaryOp, Const, DType, DTypeKind, Node, NodeId, TypedDag, UnaryOp};
use crate::opt::const_fold::const_fold;
use crate::{
    Diag, DiagCode, DiagConfig, OptConfig, Severity, Value, lower_with_config_collect,
    optimize_with_info_and_config, parse_dsl_with_errors,
};

#[derive(Debug, Clone)]
pub struct ReplConfig {
    pub opt_enabled: bool,
    pub print_simplified_expr: bool,
    pub header_locals_limit: usize,
}

impl Default for ReplConfig {
    fn default() -> Self {
        Self {
            opt_enabled: true,
            print_simplified_expr: true,
            header_locals_limit: 5,
        }
    }
}

#[derive(Debug, Clone, Default)]
pub struct ReplHeader {
    pub inputs_summary: String,
    pub locals_summary: String,
    pub last_diag_summary: String,
}

#[derive(Debug, Default)]
pub struct ReplOutput {
    pub header: ReplHeader,
    pub messages: Vec<String>,
}

#[derive(Debug, Clone)]
struct LocalBinding {
    expr_src: String,
    dtype: DType,
}

#[derive(Debug)]
pub struct ReplSession {
    inputs: Vec<(String, DType)>,
    inputs_map: HashMap<String, DType>,
    locals: HashMap<String, LocalBinding>,
    local_order: Vec<String>,
    last_diags: Vec<Diag>,
    config: ReplConfig,
    diag_cfg: DiagConfig,
    opt_cfg: OptConfig,
}

impl ReplSession {
    pub fn new(config: ReplConfig) -> Self {
        Self {
            inputs: Vec::new(),
            inputs_map: HashMap::new(),
            locals: HashMap::new(),
            local_order: Vec::new(),
            last_diags: Vec::new(),
            config,
            diag_cfg: DiagConfig::default(),
            opt_cfg: OptConfig::default(),
        }
    }

    pub fn set_diag_config(&mut self, diag_cfg: DiagConfig) {
        self.diag_cfg = diag_cfg;
    }

    pub fn set_opt_config(&mut self, opt_cfg: OptConfig) {
        self.opt_cfg = opt_cfg;
    }

    pub fn header(&self) -> ReplHeader {
        self.make_header()
    }

    pub fn process_line(&mut self, line: &str) -> ReplOutput {
        let trimmed = line.trim();
        self.last_diags.clear();

        if trimmed.is_empty() {
            return self.make_output();
        }

        match Self::classify_line(trimmed) {
            ReplLine::InputDecl { name, dtype_src } => {
                let mut out = self.make_output();
                match DType::from_str(dtype_src) {
                    Ok(dtype) => {
                        if !self.inputs_map.contains_key(name) {
                            self.inputs.push((name.to_string(), dtype));
                        } else if let Some(entry) = self.inputs.iter_mut().find(|(n, _)| n == name)
                        {
                            entry.1 = dtype;
                        }
                        self.inputs_map.insert(name.to_string(), dtype);
                        out.messages
                            .push(format!("input {name}: {}", dtype.to_str()));
                    }
                    Err(_) => {
                        let diag = Diag::new(
                            DiagCode::ParseUnknownDType,
                            format!("unknown dtype `{dtype_src}`"),
                            None,
                            None,
                        );
                        self.last_diags.push(diag.clone());
                        let sev = self.diag_cfg.severity(diag.code);
                        out.messages.push(format_diag_line(&diag, sev));
                    }
                }
                out.header = self.make_header();
                out
            }
            ReplLine::Assign { name, expr_src } => {
                let synthetic = self.build_synthetic_source(Some((name, expr_src)), None);
                self.eval_program_line(&synthetic, Some(name), Some(expr_src.to_string()), true)
            }
            ReplLine::Expr { expr_src } => {
                let synthetic = self.build_synthetic_source(None, Some(expr_src));
                self.eval_program_line(&synthetic, None, None, false)
            }
        }
    }

    fn eval_program_line(
        &mut self,
        source: &str,
        assigned_name: Option<&str>,
        assigned_expr_src: Option<String>,
        append_value_name_expr: bool,
    ) -> ReplOutput {
        let mut out = self.make_output();
        let func = match parse_dsl_with_errors(source) {
            Ok(f) => f,
            Err(diags) => {
                self.last_diags = diags;
                for d in &self.last_diags {
                    let sev = self.diag_cfg.severity(d.code);
                    out.messages.push(format_diag_line(d, sev));
                }
                out.header = self.make_header();
                return out;
            }
        };

        let lower_result = lower_with_config_collect(&func, self.diag_cfg.clone());
        for d in &lower_result.diags {
            self.last_diags.push(d.clone());
            let sev = self.diag_cfg.severity(d.code);
            out.messages.push(format_diag_line(d, sev));
        }
        for e in &lower_result.errors {
            self.last_diags.push(e.clone());
            out.messages.push(format_diag_line(e, Severity::Error));
        }
        if !lower_result.errors.is_empty() {
            out.header = self.make_header();
            return out;
        }

        // Keep the lowered DAG (with explicit-cast marks) for pretty-printing, but
        // run optimisation on a separate copy used for execution. For printing we
        // additionally run a const-fold pass to simplify literals.
        let lowered_dag = lower_result.dag;
        let display_dag = const_fold(&lowered_dag).dag;
        let mut dag = lowered_dag.clone();
        if self.config.opt_enabled {
            match optimize_with_info_and_config(&dag, &self.opt_cfg) {
                Ok(opt) => {
                    for d in &opt.diags {
                        self.last_diags.push(d.clone());
                        let sev = self.diag_cfg.severity(d.code);
                        out.messages.push(format_diag_line(d, sev));
                    }
                    dag = opt.dag;
                }
                Err(e) => {
                    out.messages.push(format!("optimization failed: {e}"));
                }
            }
        }

        if dag.outputs.is_empty() {
            out.messages.push("empty output".to_string());
            out.header = self.make_header();
            return out;
        }

        let out_id = display_dag.outputs[0];
        let out_dtype = display_dag.dtypes[out_id.0];
        let simplified_expr = render_simplified_expr(&display_dag, out_id, true);

        // In the REPL we always go through the const-folded / simplified
        // expression path, regardless of whether there are inputs. This keeps
        // the behaviour uniform: the user sees exactly one line in the form
        // `expr: dtype`, even for pure constants.
        out.messages
            .push(format!("{simplified_expr}: {}", out_dtype.to_str()));

        if let Some(name) = assigned_name {
            let expr_for_local = if append_value_name_expr {
                assigned_expr_src.unwrap_or_else(|| simplified_expr.clone())
            } else {
                simplified_expr.clone()
            };
            self.upsert_local(name.to_string(), expr_for_local, out_dtype);
        }

        out.header = self.make_header();
        out
    }

    fn upsert_local(&mut self, name: String, expr_src: String, dtype: DType) {
        if !self.locals.contains_key(&name) {
            self.local_order.push(name.clone());
        }
        self.locals.insert(name, LocalBinding { expr_src, dtype });
    }

    fn make_output(&self) -> ReplOutput {
        let header = self.make_header();
        ReplOutput {
            header,
            messages: Vec::new(),
        }
    }

    fn classify_line(line: &str) -> ReplLine<'_> {
        if let Some((name, dtype_src)) = parse_input_decl(line) {
            return ReplLine::InputDecl { name, dtype_src };
        }
        if let Some((name, expr_src)) = parse_assignment(line) {
            return ReplLine::Assign { name, expr_src };
        }
        ReplLine::Expr { expr_src: line }
    }

    fn build_synthetic_source(&self, assign: Option<(&str, &str)>, expr: Option<&str>) -> String {
        let mut src = String::new();
        if self.inputs.is_empty() {
            src.push_str("||\n");
        } else {
            let params = self
                .inputs
                .iter()
                .map(|(name, dtype)| format!("{name}: {}", dtype.to_str()))
                .collect::<Vec<_>>()
                .join(", ");
            src.push_str(&format!("|{params}|\n"));
        }
        for name in &self.local_order {
            if let Some(binding) = self.locals.get(name) {
                src.push_str(&format!("{name} = {}\n", binding.expr_src));
            }
        }
        if let Some((name, rhs)) = assign {
            src.push_str(&format!("{name} = {rhs}\n{name}\n"));
        } else if let Some(e) = expr {
            src.push_str(e);
            src.push('\n');
        }
        src
    }

    fn make_header(&self) -> ReplHeader {
        let inputs_summary = if self.inputs.is_empty() {
            "inputs: (none)".to_string()
        } else {
            let parts = self
                .inputs
                .iter()
                .map(|(name, dtype)| format!("{name}: {}", dtype.to_str()))
                .collect::<Vec<_>>();
            format!("inputs: {}", parts.join(", "))
        };

        let locals_summary = if self.local_order.is_empty() {
            "locals: (none)".to_string()
        } else {
            let n = self.config.header_locals_limit.max(1);
            let start = self.local_order.len().saturating_sub(n);
            let mut parts = Vec::new();
            for name in &self.local_order[start..] {
                if let Some(binding) = self.locals.get(name) {
                    parts.push(format!("{name}: {}", binding.dtype.to_str()));
                } else {
                    parts.push(name.clone());
                }
            }
            format!("locals: {}", parts.join(", "))
        };

        // Only summarize Info/Warning in header; errors are already shown in messages.
        let last_diag_summary = {
            let soft: Vec<_> = self
                .last_diags
                .iter()
                .filter(|d| {
                    let sev = self.diag_cfg.severity(d.code);
                    sev == Severity::Info || sev == Severity::Warning
                })
                .collect();
            if soft.is_empty() {
                "last diag: none".to_string()
            } else {
                let mut best = soft[0];
                let mut best_sev = self.diag_cfg.severity(best.code);
                for d in &soft[1..] {
                    let sev = self.diag_cfg.severity(d.code);
                    if sev > best_sev {
                        best = d;
                        best_sev = sev;
                    }
                }
                format!(
                    "last diag: {} {} ({})",
                    best.code.short_code(),
                    best.code.as_str(),
                    severity_label(best_sev)
                )
            }
        };

        ReplHeader {
            inputs_summary,
            locals_summary,
            last_diag_summary,
        }
    }
}

#[derive(Debug, Clone, Copy)]
enum ReplLine<'a> {
    InputDecl { name: &'a str, dtype_src: &'a str },
    Assign { name: &'a str, expr_src: &'a str },
    Expr { expr_src: &'a str },
}

fn parse_input_decl(line: &str) -> Option<(&str, &str)> {
    if line.contains('=') {
        return None;
    }
    let (lhs, rhs) = line.split_once(':')?;
    let name = lhs.trim();
    let dtype_src = rhs.trim();
    if is_identifier(name) && !dtype_src.is_empty() {
        Some((name, dtype_src))
    } else {
        None
    }
}

fn parse_assignment(line: &str) -> Option<(&str, &str)> {
    let mut idx = None;
    let bytes = line.as_bytes();
    for (i, b) in bytes.iter().enumerate() {
        if *b != b'=' {
            continue;
        }
        let prev = i.checked_sub(1).and_then(|j| bytes.get(j)).copied();
        let next = bytes.get(i + 1).copied();
        if prev == Some(b'=') || next == Some(b'=') {
            continue;
        }
        idx = Some(i);
        break;
    }
    let i = idx?;
    let name = line[..i].trim();
    let expr_src = line[i + 1..].trim();
    if is_identifier(name) && !expr_src.is_empty() {
        Some((name, expr_src))
    } else {
        None
    }
}

fn is_identifier(name: &str) -> bool {
    let mut chars = name.chars();
    let Some(first) = chars.next() else {
        return false;
    };
    if !(first == '_' || first.is_ascii_alphabetic()) {
        return false;
    }
    chars.all(|c| c == '_' || c.is_ascii_alphanumeric())
}

fn severity_label(sev: Severity) -> &'static str {
    match sev {
        Severity::Allow => "Allow",
        Severity::Info => "Info",
        Severity::Warning => "Warning",
        Severity::Error => "Error",
        Severity::Fatal => "Fatal",
    }
}

fn format_diag_line(diag: &Diag, severity: Severity) -> String {
    let base = format!(
        "{}: {} {}: {}",
        severity_label(severity).to_lowercase(),
        diag.code.short_code(),
        diag.code.as_str(),
        diag.message
    );
    if let Some(hint) = &diag.hint {
        format!("{base}\n  help: {hint}")
    } else {
        base
    }
}

/// Precedence for minimal parentheses (higher = binds tighter). Used when rendering.
fn bin_prec(op: BinaryOp) -> (u8, bool) {
    use BinaryOp::*;
    match op {
        Or => (1, true),
        And => (2, true),
        Eq | Ne | Lt | Le | Gt | Ge => (3, true),
        Add | Sub => (4, true),
        Mul | Div | FloorDiv | Rem => (5, true),
        Pow => (6, false), // right-assoc
    }
}

fn node_prec(dag: &TypedDag, id: NodeId) -> u8 {
    match &dag.nodes[id.0] {
        Node::Param(_) | Node::Const(_) => 255,
        Node::Unary(..) => 7,
        Node::Binary(op, ..) => bin_prec(*op).0,
        Node::IfElse(..) => 0,
        Node::Call(..) => 9,
        Node::Cast(..) => 8,
    }
}

/// Formats parameter list as `|a: i64, b: i64|`.
pub fn format_params_signature(params: &[(String, DType)]) -> String {
    let inner = params
        .iter()
        .map(|(n, d)| format!("{}: {}", n, d.to_str()))
        .collect::<Vec<_>>()
        .join(", ");
    format!("|{inner}|")
}

/// Removes redundant outer parentheses, e.g. `((1 + (-a)))` -> `(1 + (-a))`.
fn strip_redundant_outer_parens(s: &str) -> String {
    let mut t = s.trim();
    loop {
        if t.len() < 2 || !t.starts_with('(') || !t.ends_with(')') {
            break;
        }
        let mut depth = 0u32;
        let mut close_at = None;
        for (i, c) in t.bytes().enumerate() {
            match c {
                b'(' => depth += 1,
                b')' => {
                    depth = depth.saturating_sub(1);
                    if depth == 0 {
                        close_at = Some(i);
                        break;
                    }
                }
                _ => {}
            }
        }
        if close_at != Some(t.len() - 1) {
            break;
        }
        t = t[1..t.len() - 1].trim();
    }
    t.to_string()
}

/// Renders the expression at the given node with minimal parentheses.
///
/// - When `show_cast_in_expr` is false, Cast nodes are rendered as their inner only (no `@dtype`).
/// - When `show_cast_in_expr` is true, only **explicit** casts from the surface syntax
///   (`cast<dtype>(...)` or `@dtype`) are rendered as `expr@dtype`. Implicit promotion
///   casts inserted during lowering or optimisation are hidden.
pub fn render_simplified_expr(dag: &TypedDag, id: NodeId, show_cast_in_expr: bool) -> String {
    fn wrap_if(s: String, need: bool) -> String {
        if !need {
            return s;
        }
        // Avoid double-wrapping: if already wrapped in one balanced pair, don't add another.
        if s.len() >= 2 && s.starts_with('(') && s.ends_with(')') {
            let mut depth = 0u32;
            for (i, c) in s.bytes().enumerate() {
                match c {
                    b'(' => depth += 1,
                    b')' => {
                        depth = depth.saturating_sub(1);
                        if depth == 0 && i == s.len() - 1 {
                            return s;
                        }
                        if depth == 0 {
                            break;
                        }
                    }
                    _ => {}
                }
            }
        }
        format!("({s})")
    }
    fn go(
        dag: &TypedDag,
        id: NodeId,
        parent_prec: u8,
        right_child: bool,
        show_cast: bool,
        cache: &mut HashMap<NodeId, String>,
    ) -> String {
        if let Some(s) = cache.get(&id) {
            return s.clone();
        }
        let (s, self_prec) = match &dag.nodes[id.0] {
            Node::Param(i) => (
                dag.params
                    .get(i.0)
                    .map(|(n, _)| n.clone())
                    .unwrap_or_else(|| format!("arg{}", i.0)),
                255,
            ),
            Node::Const(c) => (
                match c {
                    Const::Bool(v) => v.to_string(),
                    Const::Int(v) => v.to_string(),
                    Const::Float(v) => v.to_string(),
                    Const::Complex(re, im) => format!("{re}+{im}j"),
                },
                255,
            ),
            Node::Unary(UnaryOp::Neg, a) => {
                let inner = go(dag, *a, 7, false, show_cast, cache);
                let need_paren = matches!(
                    dag.nodes[a.0],
                    Node::Binary(..) | Node::IfElse(..) | Node::Unary(..)
                );
                (format!("-{}", wrap_if(inner, need_paren)), 7)
            }
            Node::Unary(UnaryOp::Not, a) => {
                let inner = go(dag, *a, 7, false, show_cast, cache);
                let need_paren = matches!(dag.nodes[a.0], Node::Binary(..) | Node::IfElse(..));
                (format!("!{}", wrap_if(inner, need_paren)), 7)
            }
            Node::Binary(op, a, b) => {
                let (prec, _left_assoc) = bin_prec(*op);
                let sym = match op {
                    BinaryOp::Add => "+",
                    BinaryOp::Sub => "-",
                    BinaryOp::Mul => "*",
                    BinaryOp::Div => "/",
                    BinaryOp::FloorDiv => "//",
                    BinaryOp::Rem => "%",
                    BinaryOp::Pow => "**",
                    BinaryOp::Eq => "==",
                    BinaryOp::Ne => "!=",
                    BinaryOp::Lt => "<",
                    BinaryOp::Le => "<=",
                    BinaryOp::Gt => ">",
                    BinaryOp::Ge => ">=",
                    BinaryOp::And => "&&",
                    BinaryOp::Or => "||",
                };
                let left_s = go(dag, *a, prec, false, show_cast, cache);
                let right_s = go(dag, *b, prec, true, show_cast, cache);
                (format!("{} {sym} {}", left_s, right_s), prec)
            }
            Node::IfElse(c, t, e) => {
                let c_s = go(dag, *c, 0, false, show_cast, cache);
                let t_s = go(dag, *t, 0, false, show_cast, cache);
                let e_s = go(dag, *e, 0, false, show_cast, cache);
                let wrap_c = matches!(dag.nodes[c.0], Node::Binary(..) | Node::IfElse(..));
                (
                    format!(
                        "if {} then {} else {}",
                        wrap_if(c_s, wrap_c),
                        wrap_if(
                            t_s,
                            matches!(dag.nodes[t.0], Node::Binary(..) | Node::IfElse(..))
                        ),
                        wrap_if(
                            e_s,
                            matches!(dag.nodes[e.0], Node::Binary(..) | Node::IfElse(..))
                        )
                    ),
                    0,
                )
            }
            Node::Call(name, args) => {
                let arg_s = args
                    .iter()
                    .map(|a| go(dag, *a, 0, false, show_cast, cache))
                    .collect::<Vec<_>>()
                    .join(", ");
                (format!("{name}({arg_s})"), 9)
            }
            Node::Cast(dtype, a, is_explicit) => {
                let inner = go(dag, *a, 8, false, show_cast, cache);
                if show_cast && *is_explicit {
                    (format!("{}@{}", inner, dtype.to_str()), 8)
                } else {
                    let inner_prec = node_prec(dag, *a);
                    (inner, inner_prec)
                }
            }
        };
        let out = if self_prec < parent_prec || (self_prec == parent_prec && right_child) {
            wrap_if(s, true)
        } else {
            s
        };
        cache.insert(id, out.clone());
        out
    }
    let mut cache: HashMap<NodeId, String> = HashMap::new();
    strip_redundant_outer_parens(&go(dag, id, 0, false, show_cast_in_expr, &mut cache))
}

pub fn format_value_with_dtype(v: &Value, dtype: DType) -> String {
    let suffix = dtype.to_str();
    match (dtype.kind, v) {
        (DTypeKind::Bool, Value::Bool(b)) => format!("{b}: {suffix}"),
        (DTypeKind::UInt, Value::UInt { value, .. }) => format!("{value}: {suffix}"),
        (DTypeKind::SInt, Value::SInt { value, .. }) => format!("{value}: {suffix}"),
        (DTypeKind::Float(_), _) => {
            let f = match v {
                Value::F16(x) => x.to_f64(),
                Value::BF16(x) => x.to_f64(),
                Value::F32(x) => *x as f64,
                Value::F64(x) => *x,
                _ => return format!("{v:?}: {suffix}"),
            };
            format!("{f}: {suffix}")
        }
        (DTypeKind::Complex(_), Value::C64 { re, im }) => {
            let re = *re as f64;
            let im = *im as f64;
            if im >= 0.0 {
                format!("{re}+{im}j: {suffix}")
            } else {
                format!("{re}{im}j: {suffix}")
            }
        }
        (DTypeKind::Complex(_), Value::C128 { re, im }) => {
            if *im >= 0.0 {
                format!("{re}+{im}j: {suffix}")
            } else {
                format!("{re}{im}j: {suffix}")
            }
        }
        _ => format!("{v:?}: {suffix}"),
    }
}
