use std::{
    fmt,
    sync::atomic::{AtomicU32, Ordering},
};

use crate::{
    interner::Symbol,
    lexer::{
        span::{Span, Spanned},
        token::{Keyword, Literal},
    },
};
use derive_more::Display;
use itertools::Itertools;
use macro_rules_attribute::derive;

pub type Ident = Symbol;

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct NodeId(u32);

impl NodeId {
    pub const DUMMY: NodeId = NodeId(u32::MAX);
}

impl Default for NodeId {
    fn default() -> Self {
        NodeIdCounter::next()
    }
}

#[derive(Debug, Default)]
pub struct NodeIdCounter;
impl NodeIdCounter {
    fn next() -> NodeId {
        static ID: AtomicU32 = AtomicU32::new(0);
        let id = NodeId(ID.fetch_add(1, Ordering::Relaxed));
        id
    }
}

#[derive(Debug, Clone, PartialEq, crate::Node!)]
pub struct AstNode<T> {
    pub id: NodeId,
    pub span: Span,
    data: T,
}

impl<T> AstNode<T> {
    pub fn new(data: T, span: Span) -> Self {
        Self {
            id: NodeId::default(),
            span,
            data,
        }
    }

    pub fn map<F, U>(self, f: F) -> AstNode<U>
    where
        F: FnOnce(T) -> U,
    {
        AstNode::new(f(self.data), self.span)
    }

    pub fn from_vec(v: Vec<Spanned<T>>) -> Vec<Self> {
        v.into_iter().map(|s| s.into()).collect()
    }

    pub fn from_existing(data: T, span: Span, id: NodeId) -> Self {
        Self { id, span, data }
    }

    pub fn into_inner(self) -> T {
        self.data
    }
}

impl<T> std::ops::Deref for AstNode<T> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        &self.data
    }
}

pub type Node<T> = AstNode<T>;

impl<T> From<Spanned<T>> for Node<T> {
    fn from(span: Spanned<T>) -> Self {
        Self::new(span.value, span.span)
    }
}

// ─── Keywords ────────────────────────────────────────────────────────────────

#[derive(Debug, Clone, Copy, PartialEq, Display)]
#[display(rename_all = "lowercase")]
pub enum Kw {
    When,
    Wait,
    Until,
    For,
    Proc,
    Let,
    Layout,
    Group,
    Condition,
    Start,
    If,
    Else,
}

impl Keyword for Kw {
    type Tag = Kw;

    fn matches(src: &str) -> Option<Self> {
        Some(match src {
            "when" => Kw::When,
            "wait" => Kw::Wait,
            "until" => Kw::Until,
            "for" => Kw::For,
            "proc" => Kw::Proc,
            "let" => Kw::Let,
            "layout" => Kw::Layout,
            "group" => Kw::Group,
            "condition" => Kw::Condition,
            "start" => Kw::Start,
            "if" => Kw::If,
            "else" => Kw::Else,
            _ => return None,
        })
    }

    fn tag(&self) -> Self::Tag {
        *self
    }
}

#[derive(Debug, Copy, Clone, crate::Node!)]
pub enum BinOp {
    // equality
    Eq,
    Ne,
    // comparison
    Lt,
    Le,
    Gt,
    Ge,
    // logical
    Or,
    And,
    // arithmetic
    Add,
    Sub,
    Div,
    Mul,
}

#[derive(Debug, Copy, Clone, crate::Node!)]
pub enum UnaryOp {
    Neg,
    Not,
}

#[derive(Debug, Clone, crate::Node!)]
pub enum Expr {
    Literal(Literal),
    Ident(Ident),
    /// `callee(args)`  — callee is any expression (covers plain `foo()` and chained `foo()()`)
    CallExpr {
        callee: Box<Node<Expr>>,
        args: Vec<Node<Expr>>,
    },
    /// `receiver.method(args)`
    MethodCall {
        receiver: Box<Node<Expr>>,
        method: Spanned<Ident>,
        args: Vec<Node<Expr>>,
    },
    /// `receiver.field`
    Field {
        receiver: Box<Node<Expr>>,
        name: Spanned<Ident>,
    },
    Grouping(Box<Node<Expr>>),
    Binary {
        lhs: Box<Node<Expr>>,
        rhs: Box<Node<Expr>>,
        op: BinOp,
    },
    Unary {
        rhs: Box<Node<Expr>>,
        op: UnaryOp,
    },
    Block(Block),
    If {
        cond: Box<Node<Expr>>,
        then: Box<Node<Expr>>,
        else_: Option<Box<Node<Expr>>>,
    },
}

impl Expr {
    pub fn unary(op: UnaryOp, rhs: Node<Expr>) -> Self {
        Self::Unary {
            rhs: Box::new(rhs),
            op,
        }
    }

    pub fn binary(op: BinOp, lhs: Node<Expr>, rhs: Node<Expr>) -> Self {
        Self::Binary {
            lhs: Box::new(lhs),
            rhs: Box::new(rhs),
            op,
        }
    }

    pub fn grouping(expr: Node<Expr>) -> Self {
        Self::Grouping(Box::new(expr))
    }

    /// Plain free-standing call: `foo(args)` — callee is an ident expression.
    /// Kept for convenience; the parser uses `call_expr` for the general case.
    pub fn call(name: Spanned<Ident>, args: Vec<Node<Expr>>) -> Self {
        let span = name.span;
        let callee = Node::new(Expr::Ident(name.value), span);
        Self::CallExpr {
            callee: Box::new(callee),
            args,
        }
    }

    /// General call: `callee(args)` — callee is any expression.
    pub fn call_expr(callee: Node<Expr>, args: Vec<Node<Expr>>) -> Self {
        Self::CallExpr {
            callee: Box::new(callee),
            args,
        }
    }

    /// Method call: `receiver.method(args)`.
    pub fn method_call(
        receiver: Node<Expr>,
        method: Spanned<Ident>,
        args: Vec<Node<Expr>>,
    ) -> Self {
        Self::MethodCall {
            receiver: Box::new(receiver),
            method,
            args,
        }
    }

    /// Field access: `receiver.field`.
    pub fn field(receiver: Node<Expr>, field: Spanned<Ident>) -> Self {
        Self::Field {
            receiver: Box::new(receiver),
            name: field,
        }
    }

    pub fn ident(name: Ident) -> Self {
        Self::Ident(name)
    }

    pub fn literal(lit: Literal) -> Self {
        Self::Literal(lit)
    }

    pub fn block(block: Block) -> Self {
        Self::Block(block)
    }
}

#[derive(Debug, Clone, crate::Node!)]
pub enum Annotation {
    Namespace {
        namespace: Spanned<Ident>,
    },
    Field {
        namespace: Spanned<Ident>,
        name: Spanned<Ident>,
    },
    Call {
        namespace: Spanned<Ident>,
        method: Spanned<Ident>,
        args: Vec<Node<Expr>>,
    },
}

#[derive(Debug, Clone, crate::Node!)]
pub struct Annotations(pub Vec<Spanned<Annotation>>);

#[derive(Debug, Clone, crate::Node!)]
pub struct CtorStmt {
    pub ty: Spanned<Ident>,
    pub args: Vec<Node<Expr>>,
}

#[derive(Debug, Clone, crate::Node!)]
pub struct LetStmt {
    pub name: Ident,
    pub ty: Option<Ident>,
    pub value: Node<Expr>,
}

#[derive(Debug, Clone, crate::Node!)]
pub struct IfStmt {
    pub cond: Node<Expr>,
    pub then_block: Node<Block>,
    pub else_block: Option<Node<Block>>,
}

#[derive(Debug, Clone, crate::Node!)]
pub enum Stmt {
    Let(LetStmt),
    Expr(Expr),
}

#[derive(Debug, Clone, crate::Node!)]
pub struct Block {
    pub stmts: Vec<Node<Stmt>>,
    pub tail: Option<Box<Node<Stmt>>>,
}

#[derive(Debug, Clone, crate::Node!)]
pub struct Field {
    pub name: Spanned<Ident>,
    pub value: Node<CtorStmt>,
}

#[derive(Debug, Clone, crate::Node!)]
pub struct LayoutDecl {
    pub fields: Vec<Node<Field>>,
}

#[derive(Debug, Clone, crate::Node!)]
pub struct Param {
    pub name: Spanned<Ident>,
    pub ty: Spanned<Ident>,
}

#[derive(Debug, Clone, crate::Node!)]
pub struct ProcDecl {
    pub annotations: Option<Annotations>,
    pub name: Spanned<Ident>,
    pub params: Vec<Node<Param>>,
    pub ret: Option<Spanned<Ident>>,
    pub body: Node<Block>,
}

#[derive(Debug, Clone, crate::Node!)]
pub struct WhenDecl {
    pub cond: Node<Expr>,
    pub body: Node<Block>,
}

#[derive(Debug, Clone, crate::Node!)]
pub struct CondDecl {
    pub annotations: Option<Annotations>,
    pub name: Spanned<Ident>,
    pub value: Node<Expr>,
}

#[derive(Debug, Clone, crate::Node!)]
pub enum GroupItem {
    Layout(LayoutDecl),
    Proc(ProcDecl),
    When(WhenDecl),
    Cond(CondDecl),
}

#[derive(Debug, Clone, crate::Node!)]
pub struct GroupDecl {
    pub name: Spanned<Ident>,
    pub params: Vec<Node<Param>>,
    pub items: Vec<Node<GroupItem>>,
}

#[derive(Debug, Clone, crate::Node!)]
pub enum Item {
    Start(Block),
    Layout(LayoutDecl),
    Group(GroupDecl),
    Proc(ProcDecl),
}

#[derive(Debug, Clone, crate::Node!)]
pub struct Prog {
    pub items: Vec<Node<Item>>,
}

fn write_indent(f: &mut fmt::Formatter<'_>, indent: usize) -> fmt::Result {
    for _ in 0..indent {
        write!(f, "    ")?;
    }
    Ok(())
}

fn fmt_symbol_with_tag(f: &mut fmt::Formatter<'_>, _tag: &str, symbol: Symbol) -> fmt::Result {
    write!(f, "{symbol}")
}

fn fmt_ident(f: &mut fmt::Formatter<'_>, ident: Ident) -> fmt::Result {
    write!(f, "{ident}")
}

fn write_comma_separated<T: fmt::Display>(f: &mut fmt::Formatter<'_>, values: &[T]) -> fmt::Result {
    for (index, value) in values.iter().enumerate() {
        if index > 0 {
            write!(f, ", ")?;
        }
        write!(f, "{value}")?;
    }
    Ok(())
}

fn fmt_block_with_indent(f: &mut fmt::Formatter<'_>, block: &Block, indent: usize) -> fmt::Result {
    writeln!(f, "{{")?;
    for stmt in &block.stmts {
        write_indent(f, indent + 1)?;
        writeln!(f, "{stmt}")?;
    }
    if let Some(tail) = &block.tail {
        write_indent(f, indent + 1)?;
        // Print just the expression for tail, not wrapped in let
        match &tail.data {
            Stmt::Expr(expr) => {
                writeln!(f, "{expr}")?;
            }
            Stmt::Let(let_stmt) => {
                writeln!(f, "{}", tail)?;
            }
        }
    }
    write_indent(f, indent)?;
    write!(f, "}}")
}

fn fmt_layout_with_indent(
    f: &mut fmt::Formatter<'_>,
    layout: &LayoutDecl,
    indent: usize,
) -> fmt::Result {
    write_indent(f, indent)?;
    writeln!(f, "layout {{")?;
    for field in &layout.fields {
        write_indent(f, indent + 1)?;
        writeln!(f, "{field}")?;
    }
    write_indent(f, indent)?;
    write!(f, "}}")
}

fn fmt_proc_with_indent(f: &mut fmt::Formatter<'_>, proc: &ProcDecl, indent: usize) -> fmt::Result {
    if let Some(annotation) = &proc.annotations {
        write_indent(f, indent)?;
        writeln!(f, "{annotation}")?;
    }
    write_indent(f, indent)?;
    write!(f, "proc ")?;
    fmt_ident(f, proc.name.value)?;
    write!(f, "(")?;
    write_comma_separated(f, &proc.params)?;
    write!(f, ")")?;
    if let Some(ret) = &proc.ret {
        write!(f, " -> {}", ret.value)?;
    }
    write!(f, " ")?;
    fmt_block_with_indent(f, &proc.body, indent)
}

fn fmt_when_with_indent(f: &mut fmt::Formatter<'_>, when: &WhenDecl, indent: usize) -> fmt::Result {
    write_indent(f, indent)?;
    write!(f, "when {} ", when.cond)?;
    fmt_block_with_indent(f, &when.body, indent)
}

fn fmt_cond_with_indent(f: &mut fmt::Formatter<'_>, cond: &CondDecl, indent: usize) -> fmt::Result {
    if let Some(annotation) = &cond.annotations {
        write_indent(f, indent)?;
        writeln!(f, "{annotation}")?;
    }
    write_indent(f, indent)?;
    write!(f, "condition ")?;
    fmt_ident(f, cond.name.value)?;
    write!(f, " = {}", cond.value)
}

fn fmt_group_item_with_indent(
    f: &mut fmt::Formatter<'_>,
    item: &GroupItem,
    indent: usize,
) -> fmt::Result {
    match item {
        GroupItem::Layout(layout) => fmt_layout_with_indent(f, layout, indent),
        GroupItem::Proc(proc) => fmt_proc_with_indent(f, proc, indent),
        GroupItem::When(when) => fmt_when_with_indent(f, when, indent),
        GroupItem::Cond(cond) => fmt_cond_with_indent(f, cond, indent),
    }
}

fn fmt_group_with_indent(
    f: &mut fmt::Formatter<'_>,
    group: &GroupDecl,
    indent: usize,
) -> fmt::Result {
    write_indent(f, indent)?;
    write!(f, "group ")?;
    fmt_ident(f, group.name.value)?;
    write!(f, "(")?;
    write_comma_separated(f, &group.params)?;
    writeln!(f, ") {{")?;
    for (index, item) in group.items.iter().enumerate() {
        fmt_group_item_with_indent(f, item, indent + 1)?;
        if index + 1 < group.items.len() {
            writeln!(f)?;
            writeln!(f)?;
        } else {
            writeln!(f)?;
        }
    }
    write_indent(f, indent)?;
    write!(f, "}}")
}

fn fmt_item_with_indent(f: &mut fmt::Formatter<'_>, item: &Item, indent: usize) -> fmt::Result {
    match item {
        Item::Start(block) => {
            write_indent(f, indent)?;
            write!(f, "start ")?;
            fmt_block_with_indent(f, block, indent)
        }
        Item::Layout(layout) => fmt_layout_with_indent(f, layout, indent),
        Item::Group(group) => fmt_group_with_indent(f, group, indent),
        Item::Proc(proc) => fmt_proc_with_indent(f, proc, indent),
    }
}

impl fmt::Display for Annotations {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}", self.0.iter().format("\n"))
    }
}

impl<T> fmt::Display for AstNode<T>
where
    T: fmt::Display,
{
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}", self.data)
    }
}

impl fmt::Display for BinOp {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let op = match self {
            BinOp::Eq => "==",
            BinOp::Ne => "!=",
            BinOp::Lt => "<",
            BinOp::Le => "<=",
            BinOp::Gt => ">",
            BinOp::Ge => ">=",
            BinOp::Or => "||",
            BinOp::And => "&&",
            BinOp::Add => "+",
            BinOp::Sub => "-",
            BinOp::Div => "/",
            BinOp::Mul => "*",
        };
        write!(f, "{op}")
    }
}

impl fmt::Display for UnaryOp {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let op = match self {
            UnaryOp::Neg => "-",
            UnaryOp::Not => "!",
        };
        write!(f, "{op}")
    }
}

impl fmt::Display for Expr {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Expr::If { cond, then, else_ } => {
                write!(f, "if {cond} {then}")?;
                if let Some(else_) = else_ {
                    write!(f, " else {else_}")?;
                }
                Ok(())
            }
            Expr::Block(block) => block.fmt(f),
            Expr::Literal(Literal::Bool(v)) => write!(f, "{v}"),
            Expr::Literal(Literal::Int(v)) => write!(f, "{v}"),
            Expr::Literal(Literal::Float(v)) => write!(f, "{v}"),
            Expr::Literal(Literal::Char(v)) => write!(f, "'{v}'"),
            Expr::Literal(Literal::String(v)) => fmt_symbol_with_tag(f, "string", *v),
            Expr::Ident(ident) => fmt_ident(f, *ident),
            Expr::CallExpr { callee, args } => {
                write!(f, "{callee}(")?;
                write_comma_separated(f, args)?;
                write!(f, ")")
            }
            Expr::MethodCall {
                receiver,
                method,
                args,
            } => {
                write!(f, "{receiver}.")?;
                fmt_ident(f, method.value)?;
                write!(f, "(")?;
                write_comma_separated(f, args)?;
                write!(f, ")")
            }
            Expr::Field { receiver, name } => {
                write!(f, "{receiver}.")?;
                fmt_ident(f, name.value)
            }
            Expr::Grouping(expr) => write!(f, "({expr})"),
            Expr::Binary { lhs, rhs, op } => write!(f, "({lhs} {op} {rhs})"),
            Expr::Unary { rhs, op } => write!(f, "({op}{rhs})"),
        }
    }
}

impl fmt::Display for Annotation {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Annotation::Namespace { namespace } => {
                write!(f, "#[")?;
                fmt_ident(f, namespace.value)?;
                write!(f, "]")
            }
            Annotation::Field { namespace, name } => {
                write!(f, "#[")?;
                fmt_ident(f, namespace.value)?;
                write!(f, "::")?;
                fmt_ident(f, name.value)?;
                write!(f, "]")
            }
            Annotation::Call {
                namespace,
                method,
                args,
            } => {
                write!(f, "#[")?;
                fmt_ident(f, namespace.value)?;
                write!(f, "::")?;
                fmt_ident(f, method.value)?;
                write!(f, "(")?;
                write_comma_separated(f, args)?;
                write!(f, ")]")
            }
        }
    }
}

impl fmt::Display for CtorStmt {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt_ident(f, self.ty.value)?;
        write!(f, "(")?;
        write_comma_separated(f, &self.args)?;
        write!(f, ")")
    }
}

impl fmt::Display for LetStmt {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "let ")?;
        fmt_ident(f, self.name)?;
        if let Some(ty) = self.ty {
            write!(f, ": ")?;
            fmt_ident(f, ty)?;
        }
        write!(f, " = {}", self.value)
    }
}

impl fmt::Display for Stmt {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Stmt::Let(let_stmt) => write!(f, "{let_stmt}"),
            Stmt::Expr(expr) => write!(f, "{expr}"),
        }
    }
}

impl fmt::Display for IfStmt {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "if {cond} {then_block}",
            cond = self.cond,
            then_block = self.then_block
        )?;
        if let Some(else_block) = &self.else_block {
            write!(f, " else {else_block}")?;
        }
        Ok(())
    }
}

impl fmt::Display for Block {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt_block_with_indent(f, self, 0)
    }
}

impl fmt::Display for Field {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt_ident(f, self.name.value)?;
        write!(f, ": {}", self.value)
    }
}

impl fmt::Display for LayoutDecl {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt_layout_with_indent(f, self, 0)
    }
}

impl fmt::Display for Param {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt_ident(f, self.name.value)?;
        write!(f, ": ")?;
        fmt_ident(f, self.ty.value)
    }
}

impl fmt::Display for ProcDecl {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt_proc_with_indent(f, self, 0)
    }
}

impl fmt::Display for WhenDecl {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt_when_with_indent(f, self, 0)
    }
}

impl fmt::Display for CondDecl {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt_cond_with_indent(f, self, 0)
    }
}

impl fmt::Display for GroupItem {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt_group_item_with_indent(f, self, 0)
    }
}

impl fmt::Display for GroupDecl {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt_group_with_indent(f, self, 0)
    }
}

impl fmt::Display for Item {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt_item_with_indent(f, self, 0)
    }
}

impl fmt::Display for Prog {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        for (index, item) in self.items.iter().enumerate() {
            if index > 0 {
                writeln!(f)?;
                writeln!(f)?;
            }
            fmt_item_with_indent(f, item, 0)?;
        }
        Ok(())
    }
}
