use std::collections::{HashMap, HashSet};
use cord_trig::ir::{NodeId, TrigGraph, TrigOp};
use crate::token::Token;

#[derive(Clone)]
pub(crate) struct UserFunc {
    pub(crate) params: Vec<String>,
    pub(crate) defaults: Vec<Option<Vec<Token>>>,
    pub(crate) body: Vec<Token>,
}

#[derive(Clone)]
pub(crate) struct Schematic {
    pub(crate) params: Vec<String>,
    pub(crate) defaults: Vec<Option<Vec<Token>>>,
    pub(crate) body: Vec<Token>,
    pub(crate) value_returning: bool,
}

pub(crate) struct ExprParser<'a> {
    pub(crate) tokens: &'a [Token],
    pub(crate) token_lines: &'a [usize],
    pub(crate) source_lines: &'a [&'a str],
    pub(crate) pos: usize,
    pub(crate) graph: TrigGraph,
    pub(crate) input_x: Option<NodeId>,
    pub(crate) input_y: Option<NodeId>,
    pub(crate) input_z: Option<NodeId>,
    pub(crate) ref_a: Option<&'a TrigGraph>,
    pub(crate) ref_b: Option<&'a TrigGraph>,
    pub(crate) vars: HashMap<String, NodeId>,
    pub(crate) funcs: HashMap<String, UserFunc>,
    pub(crate) schematics: HashMap<String, Schematic>,
    pub(crate) objects: Vec<String>,
    pub(crate) object_nodes: HashMap<String, NodeId>,
    pub(crate) obj_results: HashSet<NodeId>,
    pub(crate) cast_nodes: Vec<(String, NodeId)>,
    pub(crate) cast_all: bool,
    pub(crate) plot_nodes: Vec<NodeId>,
    pub(crate) plot_all: bool,
    pub(crate) bare_exprs: Vec<NodeId>,
    pub(crate) vars_since_last_cast: u32,
    pub(crate) exprs_since_last_plot: u32,
    pub(crate) used_vars: HashSet<String>,
    pub(crate) warnings: Vec<String>,
}

impl<'a> ExprParser<'a> {
    pub(crate) fn new(
        tokens: &'a [Token],
        token_lines: &'a [usize],
        source_lines: &'a [&'a str],
        ref_a: Option<&'a TrigGraph>,
        ref_b: Option<&'a TrigGraph>,
    ) -> Self {
        ExprParser {
            tokens,
            token_lines,
            source_lines,
            pos: 0,
            graph: TrigGraph::new(),
            input_x: None,
            input_y: None,
            input_z: None,
            ref_a,
            ref_b,
            vars: HashMap::new(),
            funcs: HashMap::new(),
            schematics: HashMap::new(),
            objects: Vec::new(),
            object_nodes: HashMap::new(),
            obj_results: HashSet::new(),
            cast_nodes: Vec::new(),
            cast_all: false,
            plot_nodes: Vec::new(),
            plot_all: false,
            bare_exprs: Vec::new(),
            vars_since_last_cast: 0,
            exprs_since_last_plot: 0,
            used_vars: HashSet::new(),
            warnings: Vec::new(),
        }
    }

    pub(crate) fn mark_obj(&mut self, node: NodeId) -> NodeId {
        self.obj_results.insert(node);
        node
    }

    pub(crate) fn is_obj_node(&self, node: NodeId) -> bool {
        self.obj_results.contains(&node)
    }

    pub(crate) fn get_x(&mut self) -> NodeId {
        *self.input_x.get_or_insert_with(|| self.graph.push(TrigOp::InputX))
    }

    pub(crate) fn get_y(&mut self) -> NodeId {
        *self.input_y.get_or_insert_with(|| self.graph.push(TrigOp::InputY))
    }

    pub(crate) fn get_z(&mut self) -> NodeId {
        *self.input_z.get_or_insert_with(|| self.graph.push(TrigOp::InputZ))
    }

    pub(crate) fn peek(&self) -> Option<&Token> {
        self.tokens.get(self.pos)
    }

    pub(crate) fn advance(&mut self) -> Option<&Token> {
        let t = self.tokens.get(self.pos);
        self.pos += 1;
        t
    }

    pub(crate) fn skip_separators(&mut self) {
        while matches!(self.peek(), Some(Token::Semi) | Some(Token::Newline)) {
            self.advance();
        }
    }

    pub(crate) fn current_line(&self) -> usize {
        let idx = if self.pos > 0 { self.pos - 1 } else { 0 };
        self.token_lines.get(idx).copied().unwrap_or(0)
    }

    pub(crate) fn err_at(&self, msg: String) -> String {
        let ln = self.current_line();
        if ln == 0 || self.source_lines.is_empty() {
            return msg;
        }
        let src = self.source_lines.get(ln - 1).unwrap_or(&"");
        let src = src.trim();
        if src.is_empty() {
            format!("line {ln}: {msg}")
        } else {
            format!("line {ln}: {msg} | {src}")
        }
    }

    pub(crate) fn expect(&mut self, expected: &Token) -> Result<(), String> {
        let t = self.advance().cloned();
        match t {
            Some(ref t) if t == expected => Ok(()),
            Some(t) => Err(self.err_at(format!("expected {expected:?}, got {t:?}"))),
            None => Err(self.err_at(format!("expected {expected:?}, got end of input"))),
        }
    }

    pub(crate) fn parse_program(&mut self) -> Result<NodeId, String> {
        let mut last = None;
        loop {
            self.skip_separators();
            if self.pos >= self.tokens.len() { break; }

            if self.is_func_def() {
                self.parse_func_def()?;
                if self.pos >= self.tokens.len() { break; }
                continue;
            }
            if matches!(self.peek(), Some(Token::Ident(s)) if s == "sch") {
                self.parse_sch_def()?;
                if self.pos >= self.tokens.len() { break; }
                continue;
            }
            if matches!(self.peek(), Some(Token::Ident(s)) if s == "cast")
                && matches!(self.tokens.get(self.pos + 1), Some(Token::LParen))
            {
                self.advance();
                self.advance();
                if matches!(self.peek(), Some(Token::RParen)) {
                    self.advance();
                    self.cast_all = true;
                    self.vars_since_last_cast = 0;
                    self.skip_separators();
                    continue;
                }
                match self.peek().cloned() {
                    Some(Token::Ident(name)) if matches!(self.tokens.get(self.pos + 1), Some(Token::RParen)) => {
                        if !self.vars.contains_key(&name) {
                            return Err(self.err_at(format!("'{name}' is not defined")));
                        }
                        let node_id = *self.vars.get(&name).unwrap();
                        self.advance();
                        self.advance();
                        self.cast_nodes.push((name, node_id));
                        self.vars_since_last_cast = 0;
                        self.skip_separators();
                        continue;
                    }
                    _ => return Err("cast() expects a variable name or no arguments".into()),
                }
            }
            if matches!(self.peek(), Some(Token::Ident(s)) if s == "plot")
                && matches!(self.tokens.get(self.pos + 1), Some(Token::LParen))
            {
                self.advance();
                self.advance();
                if matches!(self.peek(), Some(Token::RParen)) {
                    self.advance();
                    self.plot_all = true;
                    self.exprs_since_last_plot = 0;
                    self.skip_separators();
                    continue;
                }
                let node = self.parse_additive()?;
                self.expect(&Token::RParen)?;
                self.plot_nodes.push(node);
                self.exprs_since_last_plot = 0;
                self.skip_separators();
                continue;
            }
            if let Some(Token::Ident(_)) = self.peek() {
                if matches!(self.tokens.get(self.pos + 1), Some(Token::Dot)) {
                    if let Some(Token::Ident(method)) = self.tokens.get(self.pos + 2) {
                        if method == "cast"
                           && matches!(self.tokens.get(self.pos + 3), Some(Token::LParen))
                           && matches!(self.tokens.get(self.pos + 4), Some(Token::RParen))
                        {
                            let name = match self.peek().cloned() {
                                Some(Token::Ident(n)) => n,
                                _ => unreachable!(),
                            };
                            if !self.object_nodes.contains_key(&name) {
                                if self.vars.contains_key(&name) {
                                    return Err(format!("'{name}' is not an Obj — cannot call .cast()"));
                                }
                                return Err(format!("'{name}' is not defined"));
                            }
                            let node_id = *self.object_nodes.get(&name).unwrap();
                            self.pos += 5;
                            self.cast_nodes.push((name, node_id));
                            self.vars_since_last_cast = 0;
                            self.skip_separators();
                            continue;
                        }
                    }
                }
            }
            if matches!(self.peek(), Some(Token::Ident(s)) if s == "let") {
                self.advance();
                let name = match self.advance().cloned() {
                    Some(Token::Ident(n)) => n,
                    _ => return Err(self.err_at("expected variable name after 'let'".into())),
                };
                let mut is_obj = false;
                if matches!(self.peek(), Some(Token::Colon)) {
                    self.advance();
                    match self.advance().cloned() {
                        Some(Token::Ident(ty)) => {
                            if ty == "Obj" || ty == "obj" {
                                is_obj = true;
                            }
                        }
                        _ => return Err(self.err_at("expected type name after ':'".into())),
                    }
                }
                self.expect(&Token::Eq)?;
                let val = self.parse_additive()?;
                self.vars.insert(name.clone(), val);
                if is_obj || self.is_obj_node(val) {
                    self.objects.push(name.clone());
                    self.object_nodes.insert(name, val);
                }
                self.vars_since_last_cast += 1;
                self.skip_separators();
                last = Some(val);
            } else if self.is_reassignment() {
                let name = match self.advance().cloned() {
                    Some(Token::Ident(n)) => n,
                    _ => unreachable!(),
                };
                self.advance();
                let val = self.parse_additive()?;
                if self.object_nodes.contains_key(&name) {
                    self.object_nodes.insert(name.clone(), val);
                }
                self.vars.insert(name, val);
                self.vars_since_last_cast += 1;
                self.skip_separators();
                last = Some(val);
            } else if self.pos < self.tokens.len() {
                let node = self.parse_additive()?;
                last = Some(node);
                self.bare_exprs.push(node);
                self.exprs_since_last_plot += 1;
                self.skip_separators();
            } else {
                break;
            }

            if self.pos >= self.tokens.len() {
                break;
            }
        }

        let output_node = last;
        let cast_var_names: HashSet<&str> = self.cast_nodes.iter().map(|(n, _)| n.as_str()).collect();
        let plot_nodes_set: HashSet<NodeId> = self.plot_nodes.iter().copied().collect();
        for (name, &id) in &self.vars {
            if !self.used_vars.contains(name)
                && !cast_var_names.contains(name.as_str())
                && !plot_nodes_set.contains(&id)
                && output_node != Some(id)
            {
                self.warnings.push(format!("unused variable: {name}"));
            }
        }

        match output_node {
            Some(node) => Ok(node),
            None if self.cast_all || !self.cast_nodes.is_empty()
                || self.plot_all || !self.plot_nodes.is_empty() =>
            {
                Ok(self.graph.push(TrigOp::Const(0.0)))
            }
            None => Err("empty expression".into()),
        }
    }

    fn is_reassignment(&self) -> bool {
        if let Some(Token::Ident(name)) = self.tokens.get(self.pos) {
            if self.vars.contains_key(name) {
                return matches!(self.tokens.get(self.pos + 1), Some(Token::Eq));
            }
        }
        false
    }

    pub(crate) fn parse_additive(&mut self) -> Result<NodeId, String> {
        let mut left = self.parse_multiplicative()?;
        loop {
            match self.peek() {
                Some(Token::Plus) => {
                    self.advance();
                    let right = self.parse_multiplicative()?;
                    if self.is_obj_node(left) != self.is_obj_node(right) {
                        return Err(self.err_at("cannot add Obj and Num".into()));
                    }
                    left = self.graph.push(TrigOp::Add(left, right));
                }
                Some(Token::Minus) => {
                    self.advance();
                    let right = self.parse_multiplicative()?;
                    left = self.graph.push(TrigOp::Sub(left, right));
                }
                _ => break,
            }
        }
        Ok(left)
    }

    fn parse_multiplicative(&mut self) -> Result<NodeId, String> {
        let mut left = self.parse_power()?;
        loop {
            match self.peek() {
                Some(Token::Star) => {
                    self.advance();
                    let right = self.parse_power()?;
                    left = self.graph.push(TrigOp::Mul(left, right));
                }
                Some(Token::Slash) => {
                    self.advance();
                    let right = self.parse_power()?;
                    left = self.graph.push(TrigOp::Div(left, right));
                }
                _ => break,
            }
        }
        Ok(left)
    }

    fn parse_power(&mut self) -> Result<NodeId, String> {
        let base = self.parse_unary()?;
        if self.peek() == Some(&Token::Caret) {
            self.advance();
            let exp = self.parse_unary()?;
            if let Some(TrigOp::Const(n)) = self.graph.nodes.get(exp as usize) {
                let n = *n;
                if n == 2.0 {
                    return Ok(self.graph.push(TrigOp::Mul(base, base)));
                } else if n == 3.0 {
                    let sq = self.graph.push(TrigOp::Mul(base, base));
                    return Ok(self.graph.push(TrigOp::Mul(sq, base)));
                }
            }
            Ok(self.graph.push(TrigOp::Mul(base, exp)))
        } else {
            Ok(base)
        }
    }

    fn parse_unary(&mut self) -> Result<NodeId, String> {
        if self.peek() == Some(&Token::Minus) {
            self.advance();
            let val = self.parse_unary()?;
            Ok(self.graph.push(TrigOp::Neg(val)))
        } else {
            self.parse_atom()
        }
    }

    fn parse_atom(&mut self) -> Result<NodeId, String> {
        while matches!(self.peek(), Some(Token::Newline)) {
            self.advance();
        }
        match self.advance().cloned() {
            Some(Token::Num(n)) => Ok(self.graph.push(TrigOp::Const(n))),
            Some(Token::Ident(name)) => {
                match name.as_str() {
                    "x" => Ok(self.get_x()),
                    "y" => Ok(self.get_y()),
                    "z" => Ok(self.get_z()),
                    "pi" | "PI" => Ok(self.graph.push(TrigOp::Const(std::f64::consts::PI))),
                    "tau" | "TAU" => Ok(self.graph.push(TrigOp::Const(2.0 * std::f64::consts::PI))),
                    "e" | "E" => Ok(self.graph.push(TrigOp::Const(std::f64::consts::E))),
                    "A" => {
                        if let Some(g) = self.ref_a {
                            Ok(self.inline_graph(g))
                        } else {
                            Err("A is not defined".into())
                        }
                    }
                    "B" => {
                        if let Some(g) = self.ref_b {
                            Ok(self.inline_graph(g))
                        } else {
                            Err("B is not defined".into())
                        }
                    }
                    "reg" => Ok(self.graph.push(TrigOp::Const(f64::NAN))),
                    "map" => self.parse_map(),
                    _ => {
                        if let Some(&node_id) = self.vars.get(&name) {
                            self.used_vars.insert(name);
                            return Ok(node_id);
                        }
                        self.parse_function_call(&name)
                    }
                }
            }
            Some(Token::LParen) => {
                let inner = self.parse_additive()?;
                self.expect(&Token::RParen)?;
                Ok(inner)
            }
            Some(t) => Err(format!("unexpected token: {t:?}")),
            None => Err("unexpected end of input".into()),
        }
    }

    pub(crate) fn parse_arg_list(&mut self) -> Result<Vec<NodeId>, String> {
        let mut args = Vec::new();
        self.skip_separators();
        if self.peek() == Some(&Token::RParen) {
            return Ok(args);
        }
        args.push(self.parse_additive()?);
        while { self.skip_separators(); self.peek() == Some(&Token::Comma) } {
            self.advance();
            self.skip_separators();
            args.push(self.parse_additive()?);
        }
        self.skip_separators();
        Ok(args)
    }
}

pub(crate) fn require_args(name: &str, args: &[NodeId], expected: usize) -> Result<(), String> {
    if args.len() != expected {
        Err(format!("{name}() requires {expected} argument(s), got {}", args.len()))
    } else {
        Ok(())
    }
}