use std::fmt::{Debug, Display, Formatter};

use crate::{combine::parser::Node, interner::Symbol, lexer::span::Span};

pub trait Keyword: Sized + PartialEq + Clone {
    type Tag: PartialEq + Display + Clone + Debug;
    fn matches(src: &str) -> Option<Self>;
    fn tag(&self) -> Self::Tag;
}

#[derive(Debug, Clone, PartialEq)]
pub struct Token<Kws: Keyword> {
    span: Span,
    kind: TokenKind<Kws>,
}

impl<Kws: Keyword> Node for Token<Kws> {}

impl<Kws: Keyword + Display> Display for Token<Kws> {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(f, "{} @ {}..{}", self.kind, self.span.start, self.span.end)
    }
}

impl<Kws: Keyword + Display> Display for TokenKind<Kws> {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Error => write!(f, "<error>"),
            Self::Eof => write!(f, "<eof>"),
            Self::Ident(sym) => write!(f, "{sym}"),
            Self::Keyword(kw) => write!(f, "{kw}"),
            Self::LeftParen => write!(f, "("),
            Self::RightParen => write!(f, ")"),
            Self::LeftBrace => write!(f, "{{"),
            Self::RightBrace => write!(f, "}}"),
            Self::LeftBracket => write!(f, "["),
            Self::RightBracket => write!(f, "]"),
            Self::Colon => write!(f, ":"),
            Self::Comma => write!(f, ","),
            Self::Semicolon => write!(f, ";"),
            Self::Dot => write!(f, "."),
            Self::Add => write!(f, "+"),
            Self::Sub => write!(f, "-"),
            Self::Mul => write!(f, "*"),
            Self::Div => write!(f, "/"),
            Self::Rem => write!(f, "%"),
            Self::Assign => write!(f, "="),
            Self::Eq => write!(f, "=="),
            Self::Ne => write!(f, "!="),
            Self::Lt => write!(f, "<"),
            Self::Le => write!(f, "<="),
            Self::Gt => write!(f, ">"),
            Self::Ge => write!(f, ">="),
            Self::Arrow => write!(f, "->"),
            Self::StrongArrow => write!(f, "=>"),
            Self::Not => write!(f, "!"),
            Self::And => write!(f, "&&"),
            Self::DoublePipe => write!(f, "||"),
            Self::Pipe => write!(f, "|"),
            Self::Ampersand => write!(f, "&"),
            Self::HashTag => write!(f, "#"),
            Self::Literal(lit) => write!(f, "{lit}"),
        }
    }
}

impl<Kws: Keyword> Token<Kws> {
    pub fn new<T: Into<Span>>(span: T, kind: TokenKind<Kws>) -> Self {
        Self {
            span: span.into(),
            kind,
        }
    }

    pub fn span(&self) -> Span {
        self.span
    }

    pub fn kind(&self) -> &TokenKind<Kws> {
        &self.kind
    }

    pub fn into_kind(self) -> TokenKind<Kws> {
        self.kind
    }

    pub fn is_eof(&self) -> bool {
        matches!(self.kind, TokenKind::Eof)
    }
}

#[derive(Debug, Clone, PartialEq)]
pub enum TokenKind<Kws: Keyword> {
    Error,
    Eof,
    Ident(Symbol),
    Keyword(Kws),
    // Punctuation
    LeftParen,
    RightParen,
    LeftBrace,
    RightBrace,
    LeftBracket,
    RightBracket,
    Colon,
    Comma,
    Semicolon,
    Dot,
    Arrow,
    StrongArrow,
    HashTag,
    // Operators
    Eq,
    Ne,
    Gt,
    Ge,
    Lt,
    Le,
    Add,
    Sub,
    Mul,
    Div,
    Rem,
    Not,
    // Literal
    Literal(Literal),
    Assign,
    DoublePipe,
    Pipe,
    And,
    Ampersand,
}

#[derive(Debug, Clone, PartialEq)]
pub enum TokenTag<Kws: Keyword> {
    Error,
    Eof,
    Ident,
    Keyword(Kws::Tag),
    // Punctuation
    LeftParen,
    RightParen,
    LeftBrace,
    RightBrace,
    LeftBracket,
    RightBracket,
    Colon,
    Comma,
    Semicolon,
    Dot,
    Arrow,
    StrongArrow,
    HashTag,
    // Operators
    Eq,
    Ne,
    Gt,
    Ge,
    Lt,
    Le,
    Add,
    Sub,
    Mul,
    Div,
    Rem,
    Not,
    // Literal
    Literal,
    Assign,
    DoublePipe,
    Pipe,
    And,
    Ampersand,
}

impl<Kw: Keyword> TokenKind<Kw> {
    pub fn tag(&self) -> TokenTag<Kw> {
        match self {
            TokenKind::Error => TokenTag::Error,
            TokenKind::Eof => TokenTag::Eof,
            TokenKind::Ident(_) => TokenTag::Ident,
            TokenKind::Keyword(kw) => TokenTag::Keyword(kw.tag()),
            TokenKind::LeftParen => TokenTag::LeftParen,
            TokenKind::RightParen => TokenTag::RightParen,
            TokenKind::LeftBrace => TokenTag::LeftBrace,
            TokenKind::RightBrace => TokenTag::RightBrace,
            TokenKind::LeftBracket => TokenTag::LeftBracket,
            TokenKind::RightBracket => TokenTag::RightBracket,
            TokenKind::Colon => TokenTag::Colon,
            TokenKind::Comma => TokenTag::Comma,
            TokenKind::Semicolon => TokenTag::Semicolon,
            TokenKind::Dot => TokenTag::Dot,
            TokenKind::Arrow => TokenTag::Arrow,
            TokenKind::StrongArrow => TokenTag::StrongArrow,
            TokenKind::Eq => TokenTag::Eq,
            TokenKind::Ne => TokenTag::Ne,
            TokenKind::Gt => TokenTag::Gt,
            TokenKind::Ge => TokenTag::Ge,
            TokenKind::Lt => TokenTag::Lt,
            TokenKind::Le => TokenTag::Le,
            TokenKind::Add => TokenTag::Add,
            TokenKind::Sub => TokenTag::Sub,
            TokenKind::Mul => TokenTag::Mul,
            TokenKind::Div => TokenTag::Div,
            TokenKind::Rem => TokenTag::Rem,
            TokenKind::Not => TokenTag::Not,
            TokenKind::Literal(_) => TokenTag::Literal,
            TokenKind::Assign => TokenTag::Assign,
            TokenKind::DoublePipe => TokenTag::DoublePipe,
            TokenKind::Pipe => TokenTag::Pipe,
            TokenKind::And => TokenTag::And,
            TokenKind::Ampersand => TokenTag::Ampersand,
            TokenKind::HashTag => TokenTag::HashTag,
        }
    }
}

impl<Kws: Keyword> Display for TokenTag<Kws> {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Error => write!(f, "`<error>`"),
            Self::Eof => write!(f, "`<eof>`"),
            Self::Ident => write!(f, "`<ident>`"), // use symbol display
            Self::Keyword(kw) => write!(f, "`{kw}`"),
            Self::LeftParen => write!(f, "`(`"),
            Self::RightParen => write!(f, "`)`"),
            Self::LeftBrace => write!(f, "`{{`"),
            Self::RightBrace => write!(f, "`}}`"),
            Self::LeftBracket => write!(f, "`[`"),
            Self::RightBracket => write!(f, "`]`"),
            Self::Colon => write!(f, "`:`"),
            Self::Comma => write!(f, "`,`"),
            Self::Semicolon => write!(f, "`;`"),
            Self::Dot => write!(f, "`.`"),
            Self::Arrow => write!(f, "`->`"),
            Self::StrongArrow => write!(f, "`=>`"),
            Self::Eq => write!(f, "`==`"),
            Self::Ne => write!(f, "`!=`"),
            Self::Gt => write!(f, "`>`"),
            Self::Ge => write!(f, "`>=`"),
            Self::Lt => write!(f, "`<`"),
            Self::Le => write!(f, "`<=`"),
            Self::Add => write!(f, "`+`"),
            Self::Sub => write!(f, "`-`"),
            Self::Mul => write!(f, "`*`"),
            Self::Div => write!(f, "`/`"),
            Self::Rem => write!(f, "`%`"),
            Self::Not => write!(f, "`!`"),
            Self::Assign => write!(f, "`=`"),
            Self::DoublePipe => write!(f, "`||`"),
            Self::Pipe => write!(f, "`|`"),
            Self::And => write!(f, "`&&`"),
            Self::Ampersand => write!(f, "`&`"),
            Self::Literal => write!(f, "`<literal>`"),
            Self::HashTag => write!(f, "`#`"),
        }
    }
}

#[derive(Debug, Clone, PartialEq)]
pub enum Literal {
    Bool(bool),
    Int(u64),
    Float(f64),
    Char(char),
    String(Symbol),
}

impl Display for Literal {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        match self {
            Literal::Bool(b) => write!(f, "{b}"),
            Literal::Int(i) => write!(f, "{i}"),
            Literal::Float(fl) => write!(f, "{fl}"),
            Literal::Char(c) => write!(f, "'{c}'"),
            Literal::String(s) => write!(f, "\"{s}\""),
        }
    }
}

impl Node for Literal {}
impl<Kw: Keyword> Node for TokenKind<Kw> {}
