use crate::ast::Span;

#[derive(Debug, Clone, PartialEq)]
pub enum Token {
    // Literals
    Number(f64),
    StringLit(String),
    True,
    False,

    // Identifiers and keywords
    Ident(String),
    Module,
    Function,
    If,
    Else,
    For,
    Let,

    // Symbols
    LParen,
    RParen,
    LBrace,
    RBrace,
    LBracket,
    RBracket,
    Semi,
    Comma,
    Dot,
    Assign,

    // Operators
    Plus,
    Minus,
    Star,
    Slash,
    Percent,
    Lt,
    Le,
    Gt,
    Ge,
    EqEq,
    BangEq,
    And,
    Or,
    Bang,
    Colon,
    Question,

    // Special
    Eof,
}

#[derive(Debug, Clone)]
pub struct SpannedToken {
    pub token: Token,
    pub span: Span,
}

pub struct Lexer<'a> {
    source: &'a [u8],
    pos: usize,
}

impl<'a> Lexer<'a> {
    pub fn new(source: &'a str) -> Self {
        Self {
            source: source.as_bytes(),
            pos: 0,
        }
    }

    pub fn tokenize(&mut self) -> Result<Vec<SpannedToken>, LexError> {
        let mut tokens = Vec::new();
        loop {
            self.skip_whitespace_and_comments();
            if self.pos >= self.source.len() {
                tokens.push(SpannedToken {
                    token: Token::Eof,
                    span: Span { start: self.pos, end: self.pos },
                });
                break;
            }
            tokens.push(self.next_token()?);
        }
        Ok(tokens)
    }

    fn skip_whitespace_and_comments(&mut self) {
        loop {
            // Whitespace
            while self.pos < self.source.len() && self.source[self.pos].is_ascii_whitespace() {
                self.pos += 1;
            }
            // Line comment
            if self.pos + 1 < self.source.len()
                && self.source[self.pos] == b'/'
                && self.source[self.pos + 1] == b'/'
            {
                while self.pos < self.source.len() && self.source[self.pos] != b'\n' {
                    self.pos += 1;
                }
                continue;
            }
            // Block comment
            if self.pos + 1 < self.source.len()
                && self.source[self.pos] == b'/'
                && self.source[self.pos + 1] == b'*'
            {
                self.pos += 2;
                let mut depth = 1u32;
                while self.pos + 1 < self.source.len() && depth > 0 {
                    if self.source[self.pos] == b'/' && self.source[self.pos + 1] == b'*' {
                        depth += 1;
                        self.pos += 2;
                    } else if self.source[self.pos] == b'*' && self.source[self.pos + 1] == b'/' {
                        depth -= 1;
                        self.pos += 2;
                    } else {
                        self.pos += 1;
                    }
                }
                continue;
            }
            break;
        }
    }

    fn next_token(&mut self) -> Result<SpannedToken, LexError> {
        let start = self.pos;
        let ch = self.source[self.pos];

        // Numbers
        if ch.is_ascii_digit() || (ch == b'.' && self.peek_is_digit()) {
            return self.lex_number(start);
        }

        // Strings
        if ch == b'"' {
            return self.lex_string(start);
        }

        // Identifiers and keywords
        if ch.is_ascii_alphabetic() || ch == b'_' || ch == b'$' {
            return Ok(self.lex_ident(start));
        }

        // Multi-char operators
        if self.pos + 1 < self.source.len() {
            let next = self.source[self.pos + 1];
            let two_char = match (ch, next) {
                (b'<', b'=') => Some(Token::Le),
                (b'>', b'=') => Some(Token::Ge),
                (b'=', b'=') => Some(Token::EqEq),
                (b'!', b'=') => Some(Token::BangEq),
                (b'&', b'&') => Some(Token::And),
                (b'|', b'|') => Some(Token::Or),
                _ => None,
            };
            if let Some(token) = two_char {
                self.pos += 2;
                return Ok(SpannedToken {
                    token,
                    span: Span { start, end: self.pos },
                });
            }
        }

        // Single-char tokens
        let token = match ch {
            b'(' => Token::LParen,
            b')' => Token::RParen,
            b'{' => Token::LBrace,
            b'}' => Token::RBrace,
            b'[' => Token::LBracket,
            b']' => Token::RBracket,
            b';' => Token::Semi,
            b',' => Token::Comma,
            b'.' => Token::Dot,
            b'=' => Token::Assign,
            b'+' => Token::Plus,
            b'-' => Token::Minus,
            b'*' => Token::Star,
            b'/' => Token::Slash,
            b'%' => Token::Percent,
            b'<' => Token::Lt,
            b'>' => Token::Gt,
            b'!' => Token::Bang,
            b':' => Token::Colon,
            b'?' => Token::Question,
            _ => {
                return Err(LexError {
                    pos: start,
                    msg: format!("unexpected character: {:?}", ch as char),
                });
            }
        };
        self.pos += 1;
        Ok(SpannedToken {
            token,
            span: Span { start, end: self.pos },
        })
    }

    fn peek_is_digit(&self) -> bool {
        self.pos + 1 < self.source.len() && self.source[self.pos + 1].is_ascii_digit()
    }

    fn lex_number(&mut self, start: usize) -> Result<SpannedToken, LexError> {
        while self.pos < self.source.len() && self.source[self.pos].is_ascii_digit() {
            self.pos += 1;
        }
        if self.pos < self.source.len() && self.source[self.pos] == b'.' {
            self.pos += 1;
            while self.pos < self.source.len() && self.source[self.pos].is_ascii_digit() {
                self.pos += 1;
            }
        }
        // Scientific notation
        if self.pos < self.source.len() && (self.source[self.pos] == b'e' || self.source[self.pos] == b'E') {
            self.pos += 1;
            if self.pos < self.source.len() && (self.source[self.pos] == b'+' || self.source[self.pos] == b'-') {
                self.pos += 1;
            }
            while self.pos < self.source.len() && self.source[self.pos].is_ascii_digit() {
                self.pos += 1;
            }
        }
        let text = std::str::from_utf8(&self.source[start..self.pos]).unwrap();
        let value: f64 = text.parse().map_err(|_| LexError {
            pos: start,
            msg: format!("invalid number: {text}"),
        })?;
        Ok(SpannedToken {
            token: Token::Number(value),
            span: Span { start, end: self.pos },
        })
    }

    fn lex_string(&mut self, start: usize) -> Result<SpannedToken, LexError> {
        self.pos += 1; // skip opening quote
        let mut s = String::new();
        while self.pos < self.source.len() && self.source[self.pos] != b'"' {
            if self.source[self.pos] == b'\\' {
                self.pos += 1;
                if self.pos >= self.source.len() {
                    return Err(LexError { pos: self.pos, msg: "unterminated string escape".into() });
                }
                match self.source[self.pos] {
                    b'n' => s.push('\n'),
                    b't' => s.push('\t'),
                    b'\\' => s.push('\\'),
                    b'"' => s.push('"'),
                    other => {
                        s.push('\\');
                        s.push(other as char);
                    }
                }
            } else {
                s.push(self.source[self.pos] as char);
            }
            self.pos += 1;
        }
        if self.pos >= self.source.len() {
            return Err(LexError { pos: start, msg: "unterminated string".into() });
        }
        self.pos += 1; // skip closing quote
        Ok(SpannedToken {
            token: Token::StringLit(s),
            span: Span { start, end: self.pos },
        })
    }

    fn lex_ident(&mut self, start: usize) -> SpannedToken {
        while self.pos < self.source.len()
            && (self.source[self.pos].is_ascii_alphanumeric()
                || self.source[self.pos] == b'_'
                || self.source[self.pos] == b'$')
        {
            self.pos += 1;
        }
        let text = std::str::from_utf8(&self.source[start..self.pos]).unwrap();
        let token = match text {
            "module" => Token::Module,
            "function" => Token::Function,
            "if" => Token::If,
            "else" => Token::Else,
            "for" => Token::For,
            "let" => Token::Let,
            "true" => Token::True,
            "false" => Token::False,
            _ => Token::Ident(text.to_string()),
        };
        SpannedToken {
            token,
            span: Span { start, end: self.pos },
        }
    }
}

#[derive(Debug, Clone)]
pub struct LexError {
    pub pos: usize,
    pub msg: String,
}

impl std::fmt::Display for LexError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "lex error at byte {}: {}", self.pos, self.msg)
    }
}

impl std::error::Error for LexError {}
