use colored::Colorize;

use crate::scanner::token::Number;
use crate::{Location, TokenKind};
use std::fmt;

/// Function parameter.
#[derive(Clone, Debug)]
pub struct Parameter {
    pub name: String,
    pub type_id: TypeId,
}

/// Direction of a node type port.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum PortDirection {
    In,
    Out,
    Inout,
}

impl fmt::Display for PortDirection {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            PortDirection::In => write!(f, "in"),
            PortDirection::Out => write!(f, "out"),
            PortDirection::Inout => write!(f, "inout"),
        }
    }
}

/// Single port specification in a node type declaration.
#[derive(Clone, Debug)]
pub struct PortSpec {
    pub direction: PortDirection,
    pub name: String,
    pub loc: Location,
}

/// Type identifier.
#[derive(Clone, Debug)]
pub enum TypeId {
    Array(Box<TypeId>, Option<Box<Expr>>),
    Boolean,
    Fn {
        params: Vec<TypeId>,
        return_type: Option<Box<TypeId>>,
    },
    Group,
    Node,
    Number,
    PortRef,
    Range,
    String,
}

/// Statement.
#[derive(Clone, Debug)]
pub enum Stmt {
    Assignment {
        target: Box<Expr>,
        expr: Box<Expr>,
        loc: Location,
    },
    Connect {
        from: Box<Expr>,
        targets: Vec<Expr>,
        loc: Location,
    },
    ConstDecl {
        name: String,
        type_id: TypeId,
        expr: Box<Expr>,
        loc: Location,
    },
    FnDecl {
        name: String,
        params: Vec<Parameter>,
        return_type: Option<TypeId>,
        body: Box<Expr>,
        loc: Location,
    },
    MemberDecl {
        name: String,
        expr: Box<Expr>,
        loc: Location,
    },
    PortDecl {
        name: String,
        exprs: Vec<Expr>,
        loc: Location,
    },
    NodeTypeDecl {
        name: String,
        ports: Vec<PortSpec>,
        loc: Location,
    },
    Return {
        expr: Option<Box<Expr>>,
        loc: Location,
    },
    VarDecl {
        name: String,
        mutable: bool,
        type_id: Option<TypeId>,
        expr: Option<Box<Expr>>,
        loc: Location,
    },
    Assert {
        condition: Box<Expr>,
        message: Option<Box<Expr>>,
        loc: Location,
    },
    UseDecl {
        name: String,
        loc: Location,
    },
    Expr {
        expr: Box<Expr>,
        discarded: bool,
    },
    Print {
        expr: Box<Expr>,
    },
}

/// Expression.
#[derive(Clone, Debug)]
pub enum Expr {
    Array {
        elements: Vec<Expr>,
        loc: Location,
    },
    ArrayFill {
        type_id: TypeId,
        length: Box<Expr>,
        loc: Location,
    },
    Binary {
        op: BinaryOp,
        lhs: Box<Expr>,
        rhs: Box<Expr>,
        loc: Location,
    },
    Break {
        value: Option<Box<Expr>>,
        loc: Location,
    },
    Block {
        stmts: Vec<Stmt>,
        result: Option<Box<Expr>>,
        loc: Location,
    },
    Call {
        callee: Box<Expr>,
        args: Vec<Expr>,
        loc: Location,
    },
    Closure {
        params: Vec<Parameter>,
        return_type: Option<TypeId>,
        body: Box<Expr>,
        loc: Location,
    },
    Continue {
        loc: Location,
    },
    Identifier {
        name: String,
        loc: Location,
    },
    If {
        condition: Box<Expr>,
        then_branch: Box<Expr>,
        else_branch: Option<Box<Expr>>,
        loc: Location,
    },
    Index {
        array: Box<Expr>,
        index: Box<Expr>,
        loc: Location,
    },
    Literal(Literal),
    MemberAccess {
        object: Box<Expr>,
        field: String,
        loc: Location,
    },
    MethodCall {
        object: Box<Expr>,
        method: String,
        args: Vec<Expr>,
        loc: Location,
    },
    Loop {
        body: Box<Expr>,
        loc: Location,
    },
    Group {
        stmts: Vec<Stmt>,
        loc: Location,
    },
    Node {
        type_name: String,
        loc: Location,
    },
    Range {
        start: Box<Expr>,
        end: Box<Expr>,
        inclusive: bool,
        loc: Location,
    },
    Ref {
        expr: Box<Expr>,
        loc: Location,
    },
    For {
        variable: String,
        var_loc: Location,
        iterable: Box<Expr>,
        body: Box<Expr>,
        loc: Location,
    },
    Self_ {
        loc: Location,
    },
    While {
        condition: Box<Expr>,
        body: Box<Expr>,
        loc: Location,
    },
    Unary {
        op: UnaryOp,
        factor: Box<Expr>,
        loc: Location,
    },
}

impl Expr {
    /// Get the location of this expression.
    #[must_use]
    pub fn loc(&self) -> &Location {
        match self {
            Expr::Array { loc, .. }
            | Expr::ArrayFill { loc, .. }
            | Expr::Binary { loc, .. }
            | Expr::Break { loc, .. }
            | Expr::Block { loc, .. }
            | Expr::Call { loc, .. }
            | Expr::Closure { loc, .. }
            | Expr::Continue { loc, .. }
            | Expr::For { loc, .. }
            | Expr::Identifier { loc, .. }
            | Expr::If { loc, .. }
            | Expr::Index { loc, .. }
            | Expr::Group { loc, .. }
            | Expr::MemberAccess { loc, .. }
            | Expr::MethodCall { loc, .. }
            | Expr::Loop { loc, .. }
            | Expr::Node { loc, .. }
            | Expr::Range { loc, .. }
            | Expr::Ref { loc, .. }
            | Expr::Self_ { loc }
            | Expr::While { loc, .. }
            | Expr::Unary { loc, .. } => loc,
            Expr::Literal(literal) => literal.loc(),
        }
    }
}

/// Literal value.
#[derive(Clone, Debug)]
pub enum Literal {
    Boolean(bool, Location),
    Number(Number, Location),
    String(String, Location),
}

impl Literal {
    /// Get the location of this literal.
    #[must_use]
    pub fn loc(&self) -> &Location {
        match self {
            Literal::Boolean(_, loc) | Literal::Number(_, loc) | Literal::String(_, loc) => loc,
        }
    }
}

/// Binary operation.
#[derive(Clone, Debug)]
pub enum BinaryOp {
    Add,
    Subtract,
    Multiply,
    Divide,
    Modulo,
    Eq,
    NotEq,
    Gt,
    GtEq,
    Lt,
    LtEq,
    LogicalAnd,
    LogicalOr,
}

/// Unary operation.
#[derive(Clone, Debug)]
pub enum UnaryOp {
    Plus,
    Minus,
    Not,
}

/// The abstract syntax tree (AST).
pub type Ast = Vec<Stmt>;

/// Conversion of token kind to binary operator.
impl From<TokenKind> for BinaryOp {
    fn from(kind: TokenKind) -> Self {
        match kind {
            TokenKind::Plus => BinaryOp::Add,
            TokenKind::Minus => BinaryOp::Subtract,
            TokenKind::Star => BinaryOp::Multiply,
            TokenKind::Slash => BinaryOp::Divide,
            TokenKind::Percent => BinaryOp::Modulo,
            TokenKind::Eq => BinaryOp::Eq,
            TokenKind::NotEq => BinaryOp::NotEq,
            TokenKind::Gt => BinaryOp::Gt,
            TokenKind::GtEq => BinaryOp::GtEq,
            TokenKind::Lt => BinaryOp::Lt,
            TokenKind::LtEq => BinaryOp::LtEq,

            _ => panic!("token kind cannot be converted to binary operator"),
        }
    }
}

/// Conversion of token kind to unary operator.
impl From<TokenKind> for UnaryOp {
    fn from(kind: TokenKind) -> Self {
        match kind {
            TokenKind::Bang => UnaryOp::Not,
            TokenKind::Minus => UnaryOp::Minus,
            TokenKind::Plus => UnaryOp::Plus,
            _ => panic!("token kind cannot be converted to unary operator"),
        }
    }
}

impl fmt::Display for Stmt {
    #[allow(clippy::too_many_lines)]
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Stmt::Assignment { target, expr, .. } => write!(
                f,
                "{} {{ target: {target}, val: {expr} }}",
                "Assignment".red(),
            ),
            Stmt::ConstDecl {
                name,
                type_id,
                expr,
                ..
            } => write!(
                f,
                "{} {{ name: {}, type: {}, val: {expr} }}",
                "ConstDecl".red(),
                name.yellow(),
                format!("{type_id:?}").blue()
            ),
            Stmt::VarDecl {
                name,
                mutable,
                type_id,
                expr,
                ..
            } => {
                let mut_str = if *mutable {
                    format!("{} ", "mut".purple())
                } else {
                    String::new()
                };
                let val_str = expr
                    .as_ref()
                    .map_or(String::new(), |e| format!(", val: {e}"));
                write!(
                    f,
                    "{} {{ name: {}{}{}{} }}",
                    "VarDecl".red(),
                    mut_str,
                    name.yellow(),
                    if type_id.is_none() {
                        String::new()
                    } else {
                        format!(": {}", format!("{:?}", type_id.clone().unwrap()).blue())
                    },
                    val_str
                )
            }
            Stmt::FnDecl {
                name,
                params,
                return_type,
                body,
                ..
            } => {
                let params_str = params
                    .iter()
                    .map(|p| format!("{}: {:?}", p.name.yellow(), p.type_id))
                    .collect::<Vec<_>>()
                    .join(", ");
                let ret_str = return_type.as_ref().map_or(String::new(), |t| {
                    format!(" -> {}", format!("{t:?}").blue())
                });
                write!(
                    f,
                    "{} {{ name: {}({}){}, body: {body} }}",
                    "FnDecl".red(),
                    name.yellow(),
                    params_str,
                    ret_str
                )
            }
            Stmt::NodeTypeDecl { name, ports, .. } => {
                let ports_str = ports
                    .iter()
                    .map(|p| format!("{} {}", p.direction, p.name.yellow()))
                    .collect::<Vec<_>>()
                    .join(", ");
                write!(
                    f,
                    "{} {{ name: {}, ports: [{}] }}",
                    "NodeTypeDecl".red(),
                    name.yellow(),
                    ports_str
                )
            }
            Stmt::Return { expr, .. } => {
                if let Some(e) = expr {
                    write!(f, "{} {{ val: {e} }}", "Return".red())
                } else {
                    write!(f, "{}", "Return".red())
                }
            }
            Stmt::Assert {
                condition, message, ..
            } => {
                if let Some(msg) = message {
                    write!(f, "{} {{ {condition}, {msg} }}", "Assert".red())
                } else {
                    write!(f, "{} {{ {condition} }}", "Assert".red())
                }
            }
            Stmt::Connect { from, targets, .. } => {
                write!(
                    f,
                    "{} {{ {from} -> {} }}",
                    "Connect".red(),
                    targets
                        .iter()
                        .map(ToString::to_string)
                        .collect::<Vec<_>>()
                        .join(", ")
                )
            }
            Stmt::MemberDecl { name, expr, .. } => {
                write!(
                    f,
                    "{} {{ self.{} = {expr} }}",
                    "MemberDecl".red(),
                    name.yellow()
                )
            }
            Stmt::PortDecl { name, exprs, .. } => write!(
                f,
                "{} {{ port {} = {} }}",
                "PortDecl".red(),
                name.yellow(),
                exprs
                    .iter()
                    .map(|e| format!("{e}"))
                    .collect::<Vec<_>>()
                    .join(", "),
            ),
            Stmt::UseDecl { name, .. } => {
                write!(f, "{} {{ {} }}", "UseDecl".red(), name.yellow())
            }
            Stmt::Expr { expr, .. } => write!(f, "{} {{ {expr} }}", "Expr".red()),
            Stmt::Print { expr } => write!(f, "{} {{ {expr} }}", "Print".red()),
        }
    }
}

impl fmt::Display for Expr {
    #[allow(clippy::too_many_lines)]
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Expr::Array { elements, .. } => {
                let elems_str = elements
                    .iter()
                    .map(|e| format!("{e}"))
                    .collect::<Vec<_>>()
                    .join(", ");
                write!(f, "{} {{ [{elems_str}] }}", "Array".red())
            }
            Expr::ArrayFill {
                type_id, length, ..
            } => {
                write!(f, "{} {{ [{:?}; {length}] }}", "ArrayFill".red(), type_id)
            }
            Expr::Break { value, .. } => {
                if let Some(val) = value {
                    write!(f, "{} {{ val: {val} }}", "Break".red())
                } else {
                    write!(f, "{}", "Break".red())
                }
            }
            Expr::Binary { op, lhs, rhs, .. } => write!(
                f,
                "{} {{ op: {}, lhs: {}, rhs: {} }}",
                "Binary".red(),
                format!("{op:?}").green(),
                lhs,
                rhs,
            ),
            Expr::Block { stmts, result, .. } => {
                let stmts_str = stmts
                    .iter()
                    .map(|s| format!("{s}"))
                    .collect::<Vec<_>>()
                    .join(", ");
                let result_str = result
                    .as_ref()
                    .map_or(String::new(), |e| format!(", result: {e}"));
                write!(
                    f,
                    "{} {{ stmts: [{stmts_str}]{result_str} }}",
                    "Block".red()
                )
            }
            Expr::Call { callee, args, .. } => {
                let args_str = args
                    .iter()
                    .map(|a| format!("{a}"))
                    .collect::<Vec<_>>()
                    .join(", ");
                write!(
                    f,
                    "{} {{ callee: {callee}, args: [{args_str}] }}",
                    "Call".red()
                )
            }
            Expr::Closure {
                params,
                return_type,
                body,
                ..
            } => {
                let params_str = params
                    .iter()
                    .map(|p| format!("{}: {:?}", p.name.yellow(), p.type_id))
                    .collect::<Vec<_>>()
                    .join(", ");
                let ret_str = return_type.as_ref().map_or(String::new(), |t| {
                    format!(" -> {}", format!("{t:?}").blue())
                });
                write!(
                    f,
                    "{} {{ |{params_str}|{ret_str} {body} }}",
                    "Closure".red()
                )
            }
            Expr::Continue { .. } => write!(f, "{}", "Continue".red()),
            Expr::Identifier { name, .. } => {
                write!(f, "{} {{ name: {} }}", "Identifier".red(), name.yellow())
            }
            Expr::If {
                condition,
                then_branch,
                else_branch,
                ..
            } => {
                let else_str = else_branch
                    .as_ref()
                    .map_or(String::new(), |e| format!(", else: {e}"));
                write!(
                    f,
                    "{} {{ cond: {condition}, then: {then_branch}{else_str} }}",
                    "If".red()
                )
            }
            Expr::Index { array, index, .. } => {
                write!(f, "{} {{ array: {array}, index: {index} }}", "Index".red())
            }
            Expr::Literal(n) => write!(f, "{} {{ val: {n} }}", "Literal".red()),
            Expr::MemberAccess { object, field, .. } => {
                write!(
                    f,
                    "{} {{ object: {object}, field: {} }}",
                    "MemberAccess".red(),
                    field.yellow()
                )
            }
            Expr::MethodCall {
                object,
                method,
                args,
                ..
            } => {
                let args_str = args
                    .iter()
                    .map(|a| format!("{a}"))
                    .collect::<Vec<_>>()
                    .join(", ");
                write!(
                    f,
                    "{} {{ object: {object}, method: {}, args: [{args_str}] }}",
                    "MethodCall".red(),
                    method.yellow()
                )
            }
            Expr::Loop { body, .. } => {
                write!(f, "{} {{ body: {body} }}", "Loop".red())
            }
            Expr::Group { stmts, .. } => {
                let stmts_str = stmts
                    .iter()
                    .map(|s| format!("{s}"))
                    .collect::<Vec<_>>()
                    .join(", ");
                write!(f, "{} {{ [{stmts_str}] }}", "Group".red())
            }
            Expr::Node { type_name, .. } => {
                write!(f, "{} {{ type: {} }}", "Node".red(), type_name.yellow())
            }
            Expr::Self_ { .. } => write!(f, "{}", "Self".red()),
            Expr::Range {
                start,
                end,
                inclusive,
                ..
            } => {
                let op = if *inclusive { "..=" } else { ".." };
                write!(
                    f,
                    "{} {{ start: {start}, {op}, end: {end} }}",
                    "Range".red()
                )
            }
            Expr::Ref { expr, .. } => {
                write!(f, "{} {{ expr: {expr} }}", "Ref".red())
            }
            Expr::For {
                variable,
                iterable,
                body,
                ..
            } => {
                write!(
                    f,
                    "{} {{ var: {}, iter: {iterable}, body: {body} }}",
                    "For".red(),
                    variable.yellow()
                )
            }
            Expr::While {
                condition, body, ..
            } => {
                write!(f, "{} {{ cond: {condition}, body: {body} }}", "While".red())
            }
            Expr::Unary { op, factor, .. } => write!(
                f,
                "{} {{ op: {}, factor: {} }}",
                "Unary".red(),
                format!("{op:?}").green(),
                factor,
            ),
        }
    }
}

impl fmt::Display for Literal {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Literal::Boolean(v, _) => write!(f, "{}", format!("{v}").purple()),
            Literal::Number(n, _) => write!(f, "{}", format!("{n}").yellow()),
            Literal::String(s, _) => {
                let escaped = s
                    .replace('\\', "\\\\")
                    .replace('\n', "\\n")
                    .replace('\r', "\\r")
                    .replace('\t', "\\t")
                    .replace('\"', "\\\"");
                write!(f, "{}", format!("\"{escaped}\"").bright_blue())
            }
        }
    }
}
