use crate::l1_ir::*;
use std::fmt;

#[derive(Debug, Clone, PartialEq)]
pub struct ParseError {
    pub message: String,
    pub line: usize,
    pub col: usize,
}

impl fmt::Display for ParseError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "parse error at {}:{}: {}", self.line, self.col, self.message)
    }
}

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

struct Lexer<'a> {
    input: &'a str,
    pos: usize,
    line: usize,
    col: usize,
}

#[derive(Debug, Clone, PartialEq)]
enum Token {
    Ident(String),
    Number(u64),
    LBrace,
    RBrace,
    LParen,
    RParen,
    LBracket,
    RBracket,
    Colon,
    Comma,
    Eq,
    Plus,
    Minus,
    Star,
    Arrow,    // ->
    Dot,
    Eof,
}

impl fmt::Display for Token {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Token::Ident(s) => write!(f, "'{}'", s),
            Token::Number(n) => write!(f, "{}", n),
            Token::LBrace => write!(f, "'{{'"),
            Token::RBrace => write!(f, "'}}'"),
            Token::LParen => write!(f, "'('"),
            Token::RParen => write!(f, "')'"),
            Token::LBracket => write!(f, "'['"),
            Token::RBracket => write!(f, "']'"),
            Token::Colon => write!(f, "':'"),
            Token::Comma => write!(f, "','"),
            Token::Eq => write!(f, "'='"),
            Token::Plus => write!(f, "'+'"),
            Token::Minus => write!(f, "'-'"),
            Token::Star => write!(f, "'*'"),
            Token::Arrow => write!(f, "'->'"),
            Token::Dot => write!(f, "'.'"),
            Token::Eof => write!(f, "EOF"),
        }
    }
}

impl<'a> Lexer<'a> {
    fn new(input: &'a str) -> Self {
        Lexer { input, pos: 0, line: 1, col: 1 }
    }

    fn skip_whitespace_and_comments(&mut self) {
        let bytes = self.input.as_bytes();
        while self.pos < bytes.len() {
            if bytes[self.pos] == b'\n' {
                self.pos += 1;
                self.line += 1;
                self.col = 1;
            } else if bytes[self.pos].is_ascii_whitespace() {
                self.pos += 1;
                self.col += 1;
            } else if self.pos + 1 < bytes.len() && bytes[self.pos] == b'/' && bytes[self.pos + 1] == b'/' {
                // Line comment
                while self.pos < bytes.len() && bytes[self.pos] != b'\n' {
                    self.pos += 1;
                }
            } else {
                break;
            }
        }
    }

    fn error(&self, msg: &str) -> ParseError {
        ParseError { message: msg.to_string(), line: self.line, col: self.col }
    }

    fn next_token(&mut self) -> Result<Token, ParseError> {
        self.skip_whitespace_and_comments();
        let bytes = self.input.as_bytes();

        if self.pos >= bytes.len() {
            return Ok(Token::Eof);
        }

        let ch = bytes[self.pos];

        // Identifiers and keywords
        if ch.is_ascii_alphabetic() || ch == b'_' {
            let start = self.pos;
            while self.pos < bytes.len() && (bytes[self.pos].is_ascii_alphanumeric() || bytes[self.pos] == b'_') {
                self.pos += 1;
                self.col += 1;
            }
            let word = &self.input[start..self.pos];
            return Ok(Token::Ident(word.to_string()));
        }

        // Numbers
        if ch.is_ascii_digit() {
            let start = self.pos;
            while self.pos < bytes.len() && bytes[self.pos].is_ascii_digit() {
                self.pos += 1;
                self.col += 1;
            }
            let num_str = &self.input[start..self.pos];
            let value = num_str.parse::<u64>()
                .map_err(|_| self.error(&format!("invalid number: {}", num_str)))?;
            return Ok(Token::Number(value));
        }

        self.pos += 1;
        self.col += 1;

        match ch {
            b'{' => Ok(Token::LBrace),
            b'}' => Ok(Token::RBrace),
            b'(' => Ok(Token::LParen),
            b')' => Ok(Token::RParen),
            b'[' => Ok(Token::LBracket),
            b']' => Ok(Token::RBracket),
            b':' => Ok(Token::Colon),
            b',' => Ok(Token::Comma),
            b'+' => Ok(Token::Plus),
            b'*' => Ok(Token::Star),
            b'.' => Ok(Token::Dot),
            b'=' => Ok(Token::Eq),
            b'-' => {
                if self.pos < bytes.len() && bytes[self.pos] == b'>' {
                    self.pos += 1;
                    self.col += 1;
                    Ok(Token::Arrow)
                } else {
                    Ok(Token::Minus)
                }
            }
            _ => Err(self.error(&format!("unexpected character: '{}'", ch as char))),
        }
    }

    fn peek_token(&mut self) -> Result<Token, ParseError> {
        let saved_pos = self.pos;
        let saved_line = self.line;
        let saved_col = self.col;
        let tok = self.next_token()?;
        self.pos = saved_pos;
        self.line = saved_line;
        self.col = saved_col;
        Ok(tok)
    }
}

pub struct Parser<'a> {
    lexer: Lexer<'a>,
}

impl<'a> Parser<'a> {
    pub fn new(input: &'a str) -> Self {
        Parser { lexer: Lexer::new(input) }
    }

    fn error(&self, msg: &str) -> ParseError {
        self.lexer.error(msg)
    }

    fn expect_ident(&mut self) -> Result<String, ParseError> {
        match self.lexer.next_token()? {
            Token::Ident(s) => Ok(s),
            other => Err(self.error(&format!("expected identifier, got {}", other))),
        }
    }

    fn expect_token(&mut self, expected: &Token) -> Result<(), ParseError> {
        let tok = self.lexer.next_token()?;
        if std::mem::discriminant(&tok) == std::mem::discriminant(expected) {
            Ok(())
        } else {
            Err(self.error(&format!("expected {}, got {}", expected, tok)))
        }
    }

    fn expect_number(&mut self) -> Result<u64, ParseError> {
        match self.lexer.next_token()? {
            Token::Number(n) => Ok(n),
            other => Err(self.error(&format!("expected number, got {}", other))),
        }
    }

    fn parse_field_type(&mut self) -> Result<FieldType, ParseError> {
        let name = self.expect_ident()?;
        match name.as_str() {
            "u64" => Ok(FieldType::U64),
            "i64" => Ok(FieldType::I64),
            "f64" => Ok(FieldType::F64),
            _ => Err(self.error(&format!("unknown type: {}", name))),
        }
    }

    fn parse_state_field(&mut self) -> Result<StateField, ParseError> {
        let name = self.expect_ident()?;
        self.expect_token(&Token::Colon)?;
        let ty = self.parse_field_type()?;
        self.expect_token(&Token::Eq)?;
        let init = self.expect_number()?;
        Ok(StateField { name, ty, init })
    }

    fn parse_state_block(&mut self) -> Result<Vec<StateField>, ParseError> {
        // "state" keyword already consumed
        self.expect_token(&Token::LBrace)?;
        let mut fields = Vec::new();
        loop {
            match self.lexer.peek_token()? {
                Token::RBrace => {
                    self.lexer.next_token()?;
                    break;
                }
                _ => fields.push(self.parse_state_field()?),
            }
        }
        Ok(fields)
    }

    fn parse_window_decl(&mut self) -> Result<WindowDecl, ParseError> {
        // "window" keyword already consumed
        let name = self.expect_ident()?;
        self.expect_token(&Token::Colon)?;
        self.expect_token(&Token::LParen)?;

        let mut fields = Vec::new();
        loop {
            match self.lexer.peek_token()? {
                Token::RParen => {
                    self.lexer.next_token()?;
                    break;
                }
                Token::Comma => {
                    self.lexer.next_token()?;
                }
                _ => fields.push(self.expect_ident()?),
            }
        }

        // readers(...)
        let readers_kw = self.expect_ident()?;
        if readers_kw != "readers" {
            return Err(self.error(&format!("expected 'readers', got '{}'", readers_kw)));
        }
        self.expect_token(&Token::LParen)?;
        let mut readers = Vec::new();
        loop {
            match self.lexer.peek_token()? {
                Token::RParen => {
                    self.lexer.next_token()?;
                    break;
                }
                Token::Comma => {
                    self.lexer.next_token()?;
                }
                _ => readers.push(self.expect_ident()?),
            }
        }

        Ok(WindowDecl { name, fields, readers })
    }

    fn parse_expr(&mut self) -> Result<Expr, ParseError> {
        let lhs = self.parse_expr_atom()?;

        // Check for assignment: ident = expr
        // or arithmetic: expr op expr
        match self.lexer.peek_token()? {
            Token::Eq => {
                // Assignment
                if let Expr::Var(name) = lhs {
                    self.lexer.next_token()?;
                    let rhs = self.parse_expr()?;
                    Ok(Expr::Assign(name, Box::new(rhs)))
                } else {
                    Err(self.error("left side of assignment must be an identifier"))
                }
            }
            Token::Plus | Token::Minus | Token::Star => {
                let op_tok = self.lexer.next_token()?;
                let op = match op_tok {
                    Token::Plus => ArithOp::Add,
                    Token::Minus => ArithOp::Sub,
                    Token::Star => ArithOp::Mul,
                    _ => unreachable!(),
                };
                let rhs = self.parse_expr_atom()?;
                Ok(Expr::Arith(op, Box::new(lhs), Box::new(rhs)))
            }
            _ => Ok(lhs),
        }
    }

    fn parse_expr_atom(&mut self) -> Result<Expr, ParseError> {
        match self.lexer.peek_token()? {
            Token::Number(_) => {
                let n = self.expect_number()?;
                Ok(Expr::Lit(n))
            }
            Token::Ident(_) => {
                let name = self.expect_ident()?;
                Ok(Expr::Var(name))
            }
            Token::LParen => {
                self.lexer.next_token()?;
                let expr = self.parse_expr()?;
                self.expect_token(&Token::RParen)?;
                Ok(expr)
            }
            other => Err(self.error(&format!("expected expression, got {}", other))),
        }
    }

    fn parse_handler(&mut self) -> Result<MessageHandler, ParseError> {
        // "on" keyword already consumed
        let message_type = self.expect_ident()?;
        self.expect_token(&Token::LParen)?;

        let mut args = Vec::new();
        loop {
            match self.lexer.peek_token()? {
                Token::RParen => {
                    self.lexer.next_token()?;
                    break;
                }
                Token::Comma => {
                    self.lexer.next_token()?;
                }
                _ => {
                    let arg_name = self.expect_ident()?;
                    self.expect_token(&Token::Colon)?;
                    let arg_ty = self.parse_field_type()?;
                    args.push((arg_name, arg_ty));
                }
            }
        }

        self.expect_token(&Token::LBrace)?;

        let mut exprs = Vec::new();
        loop {
            match self.lexer.peek_token()? {
                Token::RBrace => {
                    self.lexer.next_token()?;
                    break;
                }
                _ => exprs.push(self.parse_expr()?),
            }
        }

        let body = if exprs.len() == 1 {
            exprs.into_iter().next().unwrap()
        } else {
            Expr::Block(exprs)
        };

        Ok(MessageHandler { message_type, args, body })
    }

    fn parse_actor(&mut self) -> Result<ActorDecl, ParseError> {
        // "actor" keyword already consumed
        let name = self.expect_ident()?;
        self.expect_token(&Token::LBrace)?;

        let mut state = Vec::new();
        let mut windows = Vec::new();
        let mut handlers = Vec::new();

        loop {
            match self.lexer.peek_token()? {
                Token::RBrace => {
                    self.lexer.next_token()?;
                    break;
                }
                Token::Ident(ref kw) => {
                    let kw = kw.clone();
                    match kw.as_str() {
                        "state" => {
                            self.lexer.next_token()?;
                            state = self.parse_state_block()?;
                        }
                        "window" => {
                            self.lexer.next_token()?;
                            windows.push(self.parse_window_decl()?);
                        }
                        "on" => {
                            self.lexer.next_token()?;
                            handlers.push(self.parse_handler()?);
                        }
                        _ => return Err(self.error(&format!(
                            "unexpected keyword '{}' in actor body", kw
                        ))),
                    }
                }
                other => return Err(self.error(&format!(
                    "unexpected token {} in actor body", other
                ))),
            }
        }

        Ok(ActorDecl { name, state, windows, handlers })
    }

    fn parse_leaf_actions(&mut self) -> Result<Vec<LeafAction>, ParseError> {
        // "process" keyword already consumed
        self.expect_token(&Token::LBrace)?;

        let mut actions = Vec::new();
        loop {
            match self.lexer.peek_token()? {
                Token::RBrace => {
                    self.lexer.next_token()?;
                    break;
                }
                Token::Ident(ref kw) => {
                    let kw = kw.clone();
                    match kw.as_str() {
                        "forward" => {
                            self.lexer.next_token()?;
                            self.expect_token(&Token::LParen)?;
                            let target = self.expect_ident()?;
                            self.expect_token(&Token::Comma)?;
                            let msg_type = self.expect_ident()?;
                            self.expect_token(&Token::LParen)?;

                            let mut args = Vec::new();
                            loop {
                                match self.lexer.peek_token()? {
                                    Token::RParen => {
                                        self.lexer.next_token()?;
                                        break;
                                    }
                                    Token::Comma => {
                                        self.lexer.next_token()?;
                                    }
                                    _ => args.push(self.parse_expr()?),
                                }
                            }
                            self.expect_token(&Token::RParen)?; // outer close

                            actions.push(LeafAction::Forward {
                                target,
                                message_type: msg_type,
                                args,
                            });
                        }
                        "read" => {
                            self.lexer.next_token()?;
                            self.expect_token(&Token::LParen)?;
                            let actor = self.expect_ident()?;
                            self.expect_token(&Token::Dot)?;
                            let window = self.expect_ident()?;
                            self.expect_token(&Token::Dot)?;
                            let field = self.expect_ident()?;
                            self.expect_token(&Token::RParen)?;

                            actions.push(LeafAction::ReadWindow {
                                actor,
                                window,
                                field,
                            });
                        }
                        "emit" => {
                            self.lexer.next_token()?;
                            self.expect_token(&Token::LParen)?;
                            let expr = self.parse_expr()?;
                            self.expect_token(&Token::RParen)?;

                            actions.push(LeafAction::Emit(expr));
                        }
                        _ => return Err(self.error(&format!(
                            "unexpected action '{}' in process body", kw
                        ))),
                    }
                }
                other => return Err(self.error(&format!(
                    "unexpected token {} in process body", other
                ))),
            }
        }
        Ok(actions)
    }

    fn parse_leaf(&mut self) -> Result<LeafDecl, ParseError> {
        // "leaf" keyword already consumed
        let name = self.expect_ident()?;
        self.expect_token(&Token::LBrace)?;

        let mut reads = None;
        let mut actions = Vec::new();

        loop {
            match self.lexer.peek_token()? {
                Token::RBrace => {
                    self.lexer.next_token()?;
                    break;
                }
                Token::Ident(ref kw) => {
                    let kw = kw.clone();
                    match kw.as_str() {
                        "reads" => {
                            self.lexer.next_token()?;
                            let actor = self.expect_ident()?;
                            self.expect_token(&Token::Dot)?;
                            let window = self.expect_ident()?;
                            reads = Some(ReadClause { actor, window });
                        }
                        "process" => {
                            self.lexer.next_token()?;
                            actions = self.parse_leaf_actions()?;
                        }
                        _ => return Err(self.error(&format!(
                            "unexpected keyword '{}' in leaf body", kw
                        ))),
                    }
                }
                other => return Err(self.error(&format!(
                    "unexpected token {} in leaf body", other
                ))),
            }
        }

        Ok(LeafDecl { name, reads, actions })
    }

    fn parse_pipeline(&mut self) -> Result<PipelineDecl, ParseError> {
        // "pipeline" keyword already consumed
        let name = self.expect_ident()?;
        self.expect_token(&Token::LBrace)?;

        let mut stages = Vec::new();
        stages.push(self.expect_ident()?);

        loop {
            match self.lexer.peek_token()? {
                Token::Arrow => {
                    self.lexer.next_token()?;
                    stages.push(self.expect_ident()?);
                }
                Token::RBrace => {
                    self.lexer.next_token()?;
                    break;
                }
                other => return Err(self.error(&format!(
                    "expected '->' or '}}' in pipeline, got {}", other
                ))),
            }
        }

        Ok(PipelineDecl { name, stages })
    }

    fn parse_ident_list(&mut self) -> Result<Vec<String>, ParseError> {
        self.expect_token(&Token::LBracket)?;
        let mut items = Vec::new();
        loop {
            match self.lexer.peek_token()? {
                Token::RBracket => {
                    self.lexer.next_token()?;
                    break;
                }
                Token::Comma => {
                    self.lexer.next_token()?;
                }
                _ => items.push(self.expect_ident()?),
            }
        }
        Ok(items)
    }

    fn parse_core(&mut self) -> Result<CoreDecl, ParseError> {
        // "core" keyword already consumed
        let name = self.expect_ident()?;
        self.expect_token(&Token::LBrace)?;

        let mut actors = Vec::new();
        let mut leaves = Vec::new();
        let mut pipelines = Vec::new();
        let mut steps = 0u64;

        loop {
            match self.lexer.peek_token()? {
                Token::RBrace => {
                    self.lexer.next_token()?;
                    break;
                }
                Token::Ident(ref kw) => {
                    let kw = kw.clone();
                    match kw.as_str() {
                        "actors" => {
                            self.lexer.next_token()?;
                            self.expect_token(&Token::Colon)?;
                            actors = self.parse_ident_list()?;
                        }
                        "leaves" => {
                            self.lexer.next_token()?;
                            self.expect_token(&Token::Colon)?;
                            leaves = self.parse_ident_list()?;
                        }
                        "pipelines" => {
                            self.lexer.next_token()?;
                            self.expect_token(&Token::Colon)?;
                            pipelines = self.parse_ident_list()?;
                        }
                        "steps" => {
                            self.lexer.next_token()?;
                            self.expect_token(&Token::Colon)?;
                            steps = self.expect_number()?;
                        }
                        _ => return Err(self.error(&format!(
                            "unexpected field '{}' in core body", kw
                        ))),
                    }
                }
                other => return Err(self.error(&format!(
                    "unexpected token {} in core body", other
                ))),
            }
        }

        Ok(CoreDecl { name, actors, leaves, pipelines, steps })
    }

    pub fn parse_program(&mut self) -> Result<L1Program, ParseError> {
        let mut actors = Vec::new();
        let mut leaves = Vec::new();
        let mut pipelines = Vec::new();
        let mut cores = Vec::new();

        loop {
            match self.lexer.peek_token()? {
                Token::Eof => break,
                Token::Ident(ref kw) => {
                    let kw = kw.clone();
                    match kw.as_str() {
                        "actor" => {
                            self.lexer.next_token()?;
                            actors.push(self.parse_actor()?);
                        }
                        "leaf" => {
                            self.lexer.next_token()?;
                            leaves.push(self.parse_leaf()?);
                        }
                        "pipeline" => {
                            self.lexer.next_token()?;
                            pipelines.push(self.parse_pipeline()?);
                        }
                        "core" => {
                            self.lexer.next_token()?;
                            cores.push(self.parse_core()?);
                        }
                        _ => return Err(self.error(&format!(
                            "unexpected top-level keyword '{}'", kw
                        ))),
                    }
                }
                other => return Err(self.error(&format!(
                    "unexpected token {} at top level", other
                ))),
            }
        }

        Ok(L1Program { actors, leaves, pipelines, cores })
    }
}

/// Parse L1 DSL text into an L1 program.
pub fn parse(input: &str) -> Result<L1Program, ParseError> {
    Parser::new(input).parse_program()
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn parse_counter_example() {
        let input = include_str!("../examples/counter.l1");
        let program = parse(input).expect("counter.l1 should parse");

        assert_eq!(program.actors.len(), 1);
        assert_eq!(program.actors[0].name, "counter");
        assert_eq!(program.actors[0].state.len(), 1);
        assert_eq!(program.actors[0].state[0].name, "count");
        assert_eq!(program.actors[0].state[0].ty, FieldType::U64);
        assert_eq!(program.actors[0].state[0].init, 0);

        assert_eq!(program.actors[0].windows.len(), 1);
        assert_eq!(program.actors[0].windows[0].name, "count_view");
        assert_eq!(program.actors[0].windows[0].fields, vec!["count"]);
        assert_eq!(program.actors[0].windows[0].readers, vec!["display"]);

        assert_eq!(program.actors[0].handlers.len(), 1);
        assert_eq!(program.actors[0].handlers[0].message_type, "Increment");
        assert_eq!(program.actors[0].handlers[0].args.len(), 1);
        assert_eq!(program.actors[0].handlers[0].args[0].0, "amount");

        assert_eq!(program.leaves.len(), 2);
        assert_eq!(program.leaves[0].name, "ticker");
        assert_eq!(program.leaves[1].name, "display");
        assert!(program.leaves[1].reads.is_some());

        assert_eq!(program.pipelines.len(), 1);
        assert_eq!(program.pipelines[0].stages, vec!["ticker", "counter", "display"]);

        assert_eq!(program.cores.len(), 1);
        assert_eq!(program.cores[0].steps, 5);
    }

    #[test]
    fn parse_window_example() {
        let input = include_str!("../examples/window.l1");
        let program = parse(input).expect("window.l1 should parse");

        assert_eq!(program.actors.len(), 1);
        assert_eq!(program.actors[0].name, "accumulator");
        assert_eq!(program.actors[0].state[0].name, "total");
        assert_eq!(program.actors[0].state[0].init, 0);

        assert_eq!(program.actors[0].windows[0].name, "total_view");
        assert_eq!(program.actors[0].windows[0].readers, vec!["observe"]);

        assert_eq!(program.leaves.len(), 2);
        assert_eq!(program.leaves[0].name, "source");
        assert_eq!(program.leaves[1].name, "observe");
        assert!(program.leaves[1].reads.is_some());
        assert_eq!(program.leaves[1].actions.len(), 2);

        assert_eq!(program.cores[0].steps, 3);
    }

    #[test]
    fn parse_error_unexpected_top_level() {
        let input = "foobar { }";
        let err = parse(input).unwrap_err();
        assert!(err.message.contains("unexpected top-level keyword 'foobar'"));
    }

    #[test]
    fn parse_error_missing_brace() {
        let input = "actor counter state { count: u64 = 0 } }";
        let err = parse(input).unwrap_err();
        assert!(err.message.contains("expected '{'"));
    }

    #[test]
    fn parse_error_unknown_type() {
        let input = "actor a { state { x: bool = 0 } }";
        let err = parse(input).unwrap_err();
        assert!(err.message.contains("unknown type: bool"));
    }

    #[test]
    fn parse_error_malformed_pipeline() {
        let input = "pipeline p { a -> }";
        let err = parse(input).unwrap_err();
        assert!(err.message.contains("expected identifier"));
    }

    #[test]
    fn parsed_matches_hand_built_counter() {
        let input = include_str!("../examples/counter.l1");
        let parsed = parse(input).expect("should parse");
        let hand_built = crate::l1_ir::tests::build_counter_program();

        assert_eq!(parsed.actors.len(), hand_built.actors.len());
        assert_eq!(parsed.actors[0].name, hand_built.actors[0].name);
        assert_eq!(parsed.actors[0].state, hand_built.actors[0].state);
        assert_eq!(parsed.actors[0].windows, hand_built.actors[0].windows);
        assert_eq!(parsed.actors[0].handlers[0].message_type, hand_built.actors[0].handlers[0].message_type);
        assert_eq!(parsed.leaves.len(), hand_built.leaves.len());
        assert_eq!(parsed.pipelines, hand_built.pipelines);
        assert_eq!(parsed.cores, hand_built.cores);
    }

    #[test]
    fn parsed_matches_hand_built_window() {
        let input = include_str!("../examples/window.l1");
        let parsed = parse(input).expect("should parse");
        let hand_built = crate::l1_ir::tests::build_window_program();

        assert_eq!(parsed.actors[0].name, hand_built.actors[0].name);
        assert_eq!(parsed.actors[0].state, hand_built.actors[0].state);
        assert_eq!(parsed.actors[0].windows, hand_built.actors[0].windows);
        assert_eq!(parsed.leaves.len(), hand_built.leaves.len());
        assert_eq!(parsed.pipelines, hand_built.pipelines);
        assert_eq!(parsed.cores, hand_built.cores);
    }
}