use super::{node, unquote};
use crate::ast::source::*;
use crate::error::{NbclError, Result, Span};
use crate::parser::Rule;
use pest::iterators::Pair;
use std::sync::atomic::{AtomicU64, Ordering};

pub fn build_stmt(pair: Pair<Rule>) -> Result<Stmt> {
    let inner = match pair.as_rule() {
        Rule::stmt | Rule::node_stmt => pair.into_inner().next().unwrap(),
        _ => pair,
    };
    let span = Span::from_pair(&inner);

    match inner.as_rule() {
        Rule::let_stmt | Rule::const_stmt => {
            let mut ii = inner.clone().into_inner();
            let name = ii.next().unwrap().as_str().to_string();
            let next = ii.next().unwrap();

            let value = build_expr(next)?;
            if inner.as_rule() == Rule::let_stmt {
                Ok(Stmt::Let(name, value))
            } else {
                Ok(Stmt::Const(name, value))
            }
        }
        Rule::assign_stmt => {
            let span = Span::from_pair(&inner);

            // name
            let mut ii = inner.clone().into_inner();
            let name = build_expr(ii.next().unwrap())?;

            // assignment op
            let pair = ii.next().unwrap();
            let inner = pair.into_inner().next().unwrap();
            let assign_op = match inner.as_rule() {
                Rule::equal => AssignOp::Equal,
                Rule::plus_equal => AssignOp::PlusEqual,
                Rule::min_equal => AssignOp::MinEqual,
                Rule::mult_equal => AssignOp::MultEqual,
                Rule::div_equal => AssignOp::DivEqual,
                _ => {
                    return Err(NbclError::Parse {
                        message: format!(
                            "Unkown assign operator ('{}') encountered.",
                            inner.as_str().to_string()
                        ),
                        hint: Some(
                            "Replace the assign operator with one of these: '=', '+=', '-='."
                                .into(),
                        ),
                        span: Some(span.clone()),
                    });
                }
            };

            // expr
            let next = ii.next().unwrap();
            let value = build_expr(next)?;

            Ok(Stmt::Assign(name, assign_op, value, span))
        }
        Rule::for_stmt => {
            let mut ii = inner.into_inner();
            let pattern_pair = ii.next().unwrap();
            let mut patterns = Vec::new();
            for ident in pattern_pair.into_inner() {
                patterns.push(ident.as_str().to_string());
            }

            // skip 'in' keyword
            let _ = ii.next().unwrap();

            let iter_expr = build_expr(ii.next().unwrap())?;

            let block_pair = ii.next().unwrap();
            let body = build_block(block_pair)?;

            Ok(Stmt::For(patterns, iter_expr, body))
        }
        Rule::while_stmt => {
            let mut ii = inner.into_inner();
            let condition = build_expr(ii.next().unwrap())?;

            let block_pair = ii.next().unwrap();
            let body = build_block(block_pair)?;

            Ok(Stmt::While(condition, body))
        }

        Rule::return_stmt => {
            let mut ii = inner.into_inner();
            let return_type = if let Some(e_pair) = ii.next() {
                match e_pair.as_rule() {
                    Rule::node_invocation => {
                        let node_inv = node::build_node_invocation(e_pair)?;
                        Some(ReturnType::Node(node_inv))
                    }
                    Rule::expr => {
                        let expr = build_expr(e_pair)?;
                        Some(ReturnType::Expr(expr))
                    }
                    rt => {
                        return Err(NbclError::Ast {
                            message: format!("unknown return type: {:?}", rt),
                            hint: None,
                            span: Some(span),
                        });
                    }
                }
            } else {
                None
            };
            Ok(Stmt::Return(return_type, span))
        }

        Rule::expr_stmt => Ok(Stmt::Expr(build_expr(inner.into_inner().next().unwrap())?)),
        _ => Err(NbclError::Ast {
            message: format!("unknown Statement: {:?}", inner.as_rule()),
            hint: None,
            span: Some(span),
        }),
    }
}

pub fn build_expr(pair: Pair<Rule>) -> Result<Expr> {
    let span = Span::from_pair(&pair);
    match pair.as_rule() {
        Rule::expr | Rule::or_expr | Rule::and_expr | Rule::add_expr | Rule::mul_expr => {
            build_binop(pair)
        }
        Rule::cmp_expr => {
            let mut inner = pair.into_inner();
            let lhs = build_expr(inner.next().unwrap())?;
            if let Some(op_pair) = inner.next() {
                let op = op_pair.as_str().to_string();
                let rhs = build_expr(inner.next().unwrap())?;
                Ok(Expr { kind: ExprKind::Binary(Box::new(lhs), op, Box::new(rhs)), span })
            } else {
                Ok(lhs)
            }
        }
        Rule::unary_expr => {
            let mut inner = pair.into_inner();
            let first = inner.next().ok_or_else(|| NbclError::Ast {
                message: "Empty unary expression".to_string(),
                hint: None,
                span: Some(span.clone()),
            })?;

            if first.as_rule() == Rule::postfix_expr {
                build_expr(first)
            } else {
                let op = first.as_str().to_string();
                let operand_pair = inner.next().ok_or_else(|| NbclError::Ast {
                    message: format!("Expected operand after unary operator '{}'", op),
                    hint: None,
                    span: Some(span.clone()),
                })?;

                let operand = build_expr(operand_pair)?;
                Ok(Expr { kind: ExprKind::Unary(op, Box::new(operand)), span })
            }
        }
        Rule::assignable_lhs => {
            let mut inner = pair.into_inner();
            let mut res = build_expr(inner.next().unwrap())?;

            let mut it = inner.peekable();

            while let Some(suffix) = it.next() {
                res = match suffix.as_rule() {
                    Rule::accessor => {
                        let is_safe = suffix.as_str() == "?.";
                        let ident = it.next().unwrap().as_str().to_string();

                        Expr {
                            kind: ExprKind::Field(Box::new(res), ident, is_safe),
                            span: span.clone(),
                        }
                    }
                    Rule::expr => Expr {
                        kind: ExprKind::Index(Box::new(res), Box::new(build_expr(suffix)?)),
                        span: span.clone(),
                    },
                    _ => res,
                };
            }
            Ok(res)
        }
        Rule::postfix_expr => {
            let mut inner = pair.into_inner();
            let mut res = build_expr(inner.next().unwrap())?;

            let mut it = inner.peekable();

            while let Some(suffix) = it.next() {
                res = match suffix.as_rule() {
                    Rule::accessor => {
                        let is_safe = suffix.as_str() == "?.";
                        let ident = it.next().unwrap().as_str().to_string();

                        Expr {
                            kind: ExprKind::Field(Box::new(res), ident, is_safe),
                            span: span.clone(),
                        }
                    }
                    Rule::expr => Expr {
                        kind: ExprKind::Index(Box::new(res), Box::new(build_expr(suffix)?)),
                        span: span.clone(),
                    },
                    Rule::call_args => {
                        let args = suffix
                            .into_inner()
                            .map(|arg_pair| build_expr(arg_pair))
                            .collect::<Result<Vec<_>>>()?;

                        Expr { kind: ExprKind::Call(Box::new(res), args), span: span.clone() }
                    }
                    _ => res,
                };
            }
            Ok(res)
        }
        Rule::range_expr => {
            let mut inner = pair.into_inner();
            let start = build_expr(inner.next().unwrap())?;

            // The operator (.. or ..=)
            let op = inner.next().unwrap();
            let inclusive = op.as_str() == "..=";

            let end = build_expr(inner.next().unwrap())?;

            Ok(Expr { kind: ExprKind::Range(Box::new(start), Box::new(end), inclusive), span })
        }
        Rule::id_expression => build_expr(pair.into_inner().next().unwrap()),
        Rule::string_lit => {
            Ok(Expr { kind: ExprKind::Literal(Literal::Str(unquote(pair.as_str()))), span })
        }
        Rule::primary_expr => build_expr(pair.into_inner().next().unwrap()),
        Rule::literal => Ok(Expr { kind: ExprKind::Literal(build_literal(pair)?), span }),
        Rule::if_expr => Ok(Expr { kind: ExprKind::If(Box::new(build_if(pair)?)), span }),
        Rule::match_expr => Ok(Expr { kind: build_match(pair)?, span }),
        Rule::lambda_expr => Ok(Expr { kind: build_lambda(pair)?, span }),
        Rule::snake_ident => Ok(Expr { kind: ExprKind::Variable(pair.as_str().to_string()), span }),
        _ => Err(NbclError::Ast {
            message: format!("unknown expr: {:?}", pair.as_rule()),
            hint: None,
            span: Some(span),
        }),
    }
}

fn build_binop(pair: Pair<Rule>) -> Result<Expr> {
    let span = Span::from_pair(&pair);
    let mut inner = pair.into_inner();

    let mut lhs = build_expr(inner.next().unwrap())?;

    while let Some(op_pair) = inner.next() {
        let op_str = op_pair.as_str().to_string();

        let rhs_pair = inner.next().ok_or_else(|| NbclError::Ast {
            message: "expected operand after operator".to_string(),
            hint: Some(
                "An operator like '+' must be followed by a value, variable, or '('.".to_string(),
            ),
            span: Some(span.clone()),
        })?;

        let rhs = build_expr(rhs_pair)?;

        lhs = Expr {
            kind: ExprKind::Binary(Box::new(lhs), op_str, Box::new(rhs)),
            span: span.clone(),
        };
    }

    Ok(lhs)
}

fn build_literal(pair: Pair<Rule>) -> Result<Literal> {
    let inner = pair.into_inner().next().unwrap();
    match inner.as_rule() {
        Rule::int_lit => Ok(Literal::Int(inner.as_str().parse().unwrap())),
        Rule::float_lit => Ok(Literal::Float(inner.as_str().parse().unwrap())),
        Rule::bool_lit => Ok(Literal::Bool(inner.as_str() == "true")),
        Rule::string_lit => Ok(Literal::Str(unquote(inner.as_str()))),
        Rule::list_lit => {
            let mut exprs = Vec::new();
            for p in inner.into_inner() {
                exprs.push(build_expr(p)?);
            }
            Ok(Literal::List(exprs))
        }
        Rule::map_lit => {
            let mut pairs = Vec::new();
            for p in inner.into_inner() {
                let mut inner_pair = p.into_inner();

                let key = inner_pair.next().unwrap().as_str().to_string();

                // skip equal statement
                let _ = inner_pair.next().unwrap();

                let value = build_expr(inner_pair.next().unwrap())?;

                pairs.push((key, value));
            }
            Ok(Literal::Map(pairs))
        }
        _ => Ok(Literal::Null),
    }
}

fn build_if(pair: Pair<Rule>) -> Result<IfExpr> {
    let mut inner = pair.into_inner();

    // Handle main 'if'
    let condition = build_expr(inner.next().unwrap())?;
    let then_branch = build_branch(inner.next().unwrap())?;

    let mut else_ifs = Vec::new();
    let mut else_branch = None;

    // Handle remaining 'else if' and 'else'
    while let Some(next) = inner.next() {
        match next.as_rule() {
            Rule::else_if_branch => {
                let mut ei_inner = next.into_inner();
                let cond = build_expr(ei_inner.next().unwrap())?;
                let body = build_branch(ei_inner.next().unwrap())?;
                else_ifs.push((cond, body));
            }
            Rule::else_branch => {
                else_branch = Some(build_branch(next.into_inner().next().unwrap())?);
            }
            _ => unreachable!("Unexpected rule in if_expr: {:?}", next.as_rule()),
        }
    }

    Ok(IfExpr { condition, then_branch, else_ifs, else_branch })
}

fn build_branch(pair: Pair<Rule>) -> Result<(Vec<BodyItem>, Option<Expr>)> {
    let mut body = Vec::new();
    let mut final_expr = None;

    for p in pair.into_inner() {
        match p.as_rule() {
            Rule::stmt => body.push(BodyItem::Stmt(build_stmt(p)?)),
            Rule::node_invocation => body.push(BodyItem::Node(node::build_node_invocation(p)?)),
            Rule::expr => final_expr = Some(build_expr(p)?),
            _ => {}
        }
    }

    if final_expr.is_none() {
        if let Some(last_stmt) = body.last() {
            if let BodyItem::Stmt(Stmt::Expr(expr)) = last_stmt {
                final_expr = Some(expr.clone());
                body.pop();
            }
        }
    }

    Ok((body, final_expr))
}

fn build_match(pair: Pair<Rule>) -> Result<ExprKind> {
    let mut inner = pair.into_inner();

    // expression
    let condition = build_expr(inner.next().unwrap())?;

    // arms
    let mut arms = Vec::new();
    for arm_pair in inner {
        arms.push(build_match_arm(arm_pair)?);
    }

    Ok(ExprKind::Match(Box::new(condition), arms))
}

fn build_match_arm(pair: Pair<Rule>) -> Result<MatchArm> {
    let mut inner = pair.into_inner();

    let pattern_pair = inner.next().unwrap();
    let pattern_inner = pattern_pair.into_inner().next().unwrap();
    let is_var = match pattern_inner.as_rule() {
        Rule::snake_ident => true,
        _ => false,
    };
    let pattern = unquote(pattern_inner.as_str());

    let body_pair = inner.next().unwrap();
    let body = match body_pair.as_rule() {
        Rule::block_body => {
            let span = Span::from_pair(&body_pair);
            let block = build_block(body_pair)?;
            let mut stmts_found = Vec::new();
            for item in block.body {
                match item {
                    BodyItem::Stmt(stmt) => stmts_found.push(stmt),
                    BodyItem::Node(_) => {
                        return Err(NbclError::Ast {
                            message: "Cannot match to a node in match block body".to_string(),
                            hint: None,
                            span: Some(span),
                        });
                    }
                }
            }

            MatchBody::Block(stmts_found, block.terminator)
        }
        _ => {
            let expr = build_expr(body_pair)?;
            MatchBody::Expr(expr)
        }
    };

    Ok(MatchArm { pattern, body, is_var })
}

pub fn generate_anon_fn_name() -> String {
    static COUNTER: AtomicU64 = AtomicU64::new(0);
    let id = COUNTER.fetch_add(1, Ordering::Relaxed);
    format!("<lambda:{}>", id)
}

fn build_lambda(pair: Pair<Rule>) -> Result<ExprKind> {
    let span = Span::from_pair(&pair);
    let mut inner = pair.clone().into_inner();
    let mut params = Vec::new();

    while let Some(next) = inner.peek() {
        if next.as_rule() == Rule::lambda_param {
            let mut param_inner = inner.next().unwrap().into_inner();
            let name = param_inner.next().unwrap().as_str().to_string();
            params.push(name);
        } else {
            break;
        }
    }

    let body_pair = inner.next().ok_or_else(|| NbclError::Ast {
        message: "Lambda must have a body".into(),
        hint: None,
        span: Some(span.clone()),
    })?;

    let mut body_items = Vec::new();

    let actual_body = if body_pair.as_rule() == Rule::lambda_body {
        body_pair.into_inner().next().unwrap()
    } else {
        body_pair
    };

    match actual_body.as_rule() {
        Rule::fn_body => {
            let item = actual_body.into_inner().next().unwrap();
            match item.as_rule() {
                Rule::fn_item => {
                    let child = item.into_inner().next().unwrap();
                    match child.as_rule() {
                        Rule::node_invocation => {
                            body_items.push(BodyItem::Node(node::build_node_invocation(child)?));
                        }
                        Rule::stmt => {
                            body_items.push(BodyItem::Stmt(build_stmt(child)?));
                        }
                        _ => {}
                    }
                }
                _ => {}
            }
        }
        _ => {
            // Single expression lambda: |x| x + 1
            let expr = build_expr(actual_body)?;
            body_items.push(BodyItem::Stmt(Stmt::Expr(expr)));
        }
    }

    let fn_def = FnDef { name: generate_anon_fn_name(), params, body: body_items, span };

    Ok(ExprKind::Lambda(fn_def))
}

pub fn build_block(pair: Pair<Rule>) -> Result<Block> {
    let mut body = Vec::new();
    let mut terminator = None;

    // inner will be the contents of the { ... }
    for inner_pair in pair.into_inner() {
        match inner_pair.as_rule() {
            Rule::stmt => {
                body.push(BodyItem::Stmt(build_stmt(inner_pair)?));
            }
            Rule::node_invocation => {
                body.push(BodyItem::Node(node::build_node_invocation(inner_pair)?));
            }
            Rule::expr => {
                // In the grammar: { stmt* ~ expr? }
                // If we hit an expr, it's the implicit return/terminator
                terminator = Some(build_expr(inner_pair)?);
            }
            _ => {}
        }
    }

    Ok(Block { body, terminator })
}
