use crate::ast::*;

pub struct Parser<'a> {
    src: &'a [u8],
    pos: usize,
    pub module: ModuleSoA,
}

impl<'a> Parser<'a> {
    pub fn new(src: &'a str) -> Self {
        Self {
            src: src.as_bytes(),
            pos: 0,
            module: ModuleSoA::new(),
        }
    }

    #[inline]
    fn peek(&self) -> Option<u8> {
        self.src.get(self.pos).copied()
    }

    #[inline]
    fn advance(&mut self) {
        self.pos += 1;
    }

    #[inline]
    fn consume_if(&mut self, expected: u8) -> bool {
        if self.peek() == Some(expected) {
            self.pos += 1;
            true
        } else {
            false
        }
    }

    fn skip_whitespace(&mut self) {
        while let Some(c) = self.peek() {
            if c.is_ascii_whitespace() {
                self.advance();
            } else if c == b'/' && self.src.get(self.pos + 1) == Some(&b'/') {
                while let Some(cc) = self.peek() {
                    if cc == b'\n' {
                        break;
                    }
                    self.advance();
                }
            } else {
                break;
            }
        }
    }

    fn parse_ident(&mut self) -> String {
        let start = self.pos;
        while let Some(c) = self.peek() {
            if c.is_ascii_alphanumeric() || c == b'-' || c == b'_' || c == b'.' {
                self.advance();
            } else {
                break;
            }
        }
        std::str::from_utf8(&self.src[start..self.pos])
            .unwrap()
            .to_string()
    }

    fn parse_rhei_expr(&mut self) -> Result<String, String> {
        self.skip_whitespace();
        if self.consume_if(b'{') {
            let start = self.pos;
            let mut depth = 1;
            while depth > 0 {
                let c = self.peek().ok_or("Unexpected EOF in Rhei block")?;
                self.advance();
                if c == b'{' {
                    depth += 1;
                }
                if c == b'}' {
                    depth -= 1;
                }
            }
            let expr = std::str::from_utf8(&self.src[start..self.pos - 1])
                .unwrap()
                .trim()
                .to_string();
            Ok(expr)
        } else {
            let start = self.pos;
            let mut parens = 0;
            let mut in_str = false;
            let mut escape = false;
            while let Some(c) = self.peek() {
                if escape {
                    escape = false;
                    self.advance();
                    continue;
                }
                if c == b'\\' {
                    escape = true;
                    self.advance();
                    continue;
                }
                if c == b'"' {
                    in_str = !in_str;
                }
                if !in_str {
                    if c == b'(' {
                        parens += 1;
                    }
                    if c == b')' {
                        if parens == 0 {
                            break;
                        }
                        parens -= 1;
                    }
                    if parens == 0 && (c == b',' || c == b'\n' || c == b'}' || c == b'{') {
                        break;
                    }
                }
                self.advance();
            }
            let expr = std::str::from_utf8(&self.src[start..self.pos])
                .unwrap()
                .trim()
                .to_string();
            Ok(expr)
        }
    }

    fn parse_value(&mut self) -> Result<Value, String> {
        self.skip_whitespace();
        let c = self.peek().ok_or("Expected value, found EOF")?;
        match c {
            b'"' => {
                self.advance();
                let start = self.pos;
                while self.peek() != Some(b'"') {
                    if self.peek().is_none() {
                        return Err("Unexpected EOF inside string literal".into());
                    }
                    self.advance();
                }
                let val = std::str::from_utf8(&self.src[start..self.pos])
                    .unwrap()
                    .to_string();
                self.advance();
                Ok(Value::String(val))
            }
            b'#' => {
                self.advance();
                let mut color = String::from("#");
                color.push_str(&self.parse_ident());
                Ok(Value::Color(color))
            }
            b'$' => {
                self.advance();
                Ok(Value::Variable(self.parse_ident()))
            }
            b'!' => {
                self.advance();
                let ident = self.parse_ident();
                if ident != "rhei" {
                    return Err("Expected !rhei".into());
                }
                self.consume_if(b':');
                Ok(Value::Rhei(self.parse_rhei_expr()?))
            }
            b'0'..=b'9' | b'-' => {
                let s = self.parse_ident();
                if s.contains('.') {
                    Ok(Value::Float(s.parse().unwrap()))
                } else {
                    Ok(Value::Int(s.parse().unwrap()))
                }
            }
            _ => {
                let s = self.parse_ident();
                match s.as_str() {
                    "true" => Ok(Value::Bool(true)),
                    "false" => Ok(Value::Bool(false)),
                    "fs" => {
                        self.consume_if(b':');
                        let start = self.pos;
                        while let Some(cc) = self.peek() {
                            if cc.is_ascii_whitespace() || cc == b',' || cc == b')' || cc == b'}' {
                                break;
                            }
                            self.advance();
                        }
                        let path = std::str::from_utf8(&self.src[start..self.pos])
                            .unwrap()
                            .to_string();
                        Ok(Value::FsPath(path))
                    }
                    _ => Err(format!("Unknown value token: {}", s)),
                }
            }
        }
    }

    fn parse_properties(&mut self) -> Result<(u32, u32), String> {
        let start_idx = self.module.prop_keys.len() as u32;
        self.advance();
        loop {
            self.skip_whitespace();
            if self.consume_if(b')') {
                break;
            }
            let key = self.parse_ident();
            if key.is_empty() {
                return Err(format!(
                    "Expected property identifier, found {:?}",
                    self.peek().map(|b| b as char)
                ));
            }
            self.skip_whitespace();
            if self.consume_if(b'=') {
                let val = self.parse_value()?;
                self.module.prop_keys.push(key);
                self.module.prop_values.push(val);
            }
            self.skip_whitespace();
            self.consume_if(b',');
        }
        let len = (self.module.prop_keys.len() as u32) - start_idx;
        Ok((start_idx, len))
    }

    fn parse_block(&mut self) -> Result<(u32, u32), String> {
        self.skip_whitespace();
        if !self.consume_if(b'{') {
            return Ok((self.module.hierarchy.len() as u32, 0));
        }

        let mut direct_children = Vec::new();
        loop {
            self.skip_whitespace();
            if self.consume_if(b'}') {
                break;
            }
            let node = self.parse_node()?;
            direct_children.push(node);
        }

        let start_idx = self.module.hierarchy.len() as u32;
        let len = direct_children.len() as u32;
        self.module.hierarchy.extend(direct_children);
        Ok((start_idx, len))
    }

    fn parse_directive(&mut self) -> Result<NodeId, String> {
        self.advance();
        let name = self.parse_ident();
        let dir = match name.as_str() {
            "version" => {
                self.skip_whitespace();
                let v = self.parse_ident().parse().unwrap();
                Directive::Version(v)
            }
            "style" => Directive::Style(match self.parse_value()? {
                Value::String(s) => s,
                _ => return Err("Style must be a string".into()),
            }),
            "global" => {
                self.skip_whitespace();
                if !self.consume_if(b'$') {
                    return Err("Expected $ var".into());
                }
                let var_name = self.parse_ident();
                self.skip_whitespace();
                self.consume_if(b'=');
                let val = self.parse_value()?;
                Directive::Global {
                    name: var_name,
                    value: val,
                }
            }
            "singleton" => {
                self.skip_whitespace();
                let sname = self.parse_ident();
                self.skip_whitespace();
                let start_idx = self.module.prop_keys.len() as u32;
                if self.consume_if(b'{') {
                    loop {
                        self.skip_whitespace();
                        if self.consume_if(b'}') {
                            break;
                        }

                        let key = self.parse_ident();
                        if key.is_empty() {
                            return Err(format!(
                                "Expected singleton property identifier, found {:?}",
                                self.peek().map(|b| b as char)
                            ));
                        }
                        self.skip_whitespace();
                        self.consume_if(b'=');
                        let val = self.parse_value()?;
                        self.module.prop_keys.push(key);
                        self.module.prop_values.push(val);
                    }
                }
                let len = (self.module.prop_keys.len() as u32) - start_idx;
                Directive::Singleton {
                    name: sname,
                    prop_span: (start_idx, len),
                }
            }
            "component" => {
                self.skip_whitespace();
                let cname = self.parse_ident();
                let mut params = Vec::new();
                self.skip_whitespace();
                if self.consume_if(b'(') {
                    loop {
                        self.skip_whitespace();
                        if self.consume_if(b')') {
                            break;
                        }
                        let pname = self.parse_ident();
                        self.skip_whitespace();
                        self.consume_if(b':');
                        self.skip_whitespace();
                        let ptype = self.parse_ident();
                        params.push((pname, ptype));
                        self.skip_whitespace();
                        self.consume_if(b',');
                    }
                }
                let child_span = self.parse_block()?;
                Directive::Component {
                    name: cname,
                    params,
                    child_span,
                }
            }
            "let" => {
                self.skip_whitespace();
                self.consume_if(b'$');
                let var_name = self.parse_ident();
                self.skip_whitespace();
                self.consume_if(b'=');
                let val = self.parse_value()?;
                Directive::Let {
                    name: var_name,
                    value: val,
                }
            }
            "if" => {
                let condition = self.parse_value()?;
                let child_span = self.parse_block()?;
                let mut else_span = None;
                self.skip_whitespace();

                let backup = self.pos;
                if self.consume_if(b'@') {
                    if self.parse_ident() == "else" {
                        else_span = Some(self.parse_block()?);
                    } else {
                        self.pos = backup;
                    }
                }
                Directive::If {
                    condition,
                    child_span,
                    else_span,
                }
            }
            "each" => {
                self.skip_whitespace();
                self.consume_if(b'$');
                let item = self.parse_ident();
                self.skip_whitespace();
                let in_kw = self.parse_ident();
                if in_kw != "in" {
                    return Err("Expected 'in' in @each".into());
                }
                let collection = self.parse_value()?;
                let child_span = self.parse_block()?;
                Directive::Each {
                    item,
                    collection,
                    child_span,
                }
            }
            "on" => {
                self.skip_whitespace();
                let event = self.parse_ident();
                let mut args = Vec::new();
                self.skip_whitespace();
                if self.consume_if(b'(') {
                    loop {
                        self.skip_whitespace();
                        if self.consume_if(b')') {
                            break;
                        }
                        let k = self.parse_ident();
                        self.skip_whitespace();
                        self.consume_if(b'=');
                        let v = self.parse_value()?;
                        args.push((k, v));
                        self.skip_whitespace();
                        self.consume_if(b',');
                    }
                }
                let child_span = self.parse_block()?;
                Directive::On {
                    event,
                    args,
                    child_span,
                }
            }
            _ => return Err(format!("Unknown directive: @{}", name)),
        };
        Ok(NodeId::Directive(self.module.push_directive(dir)))
    }

    fn parse_element(&mut self) -> Result<NodeId, String> {
        let name = self.parse_ident();
        let el_id = self.module.push_element(name);

        self.skip_whitespace();
        if self.peek() == Some(b'(') {
            let span = self.parse_properties()?;
            self.module.elem_prop_spans[el_id as usize] = span;
        }

        self.skip_whitespace();
        let p = self.peek();
        if p == Some(b'{') {
            let span = self.parse_block()?;
            self.module.elem_child_spans[el_id as usize] = span;
        } else if p == Some(b'"') || p == Some(b'$') {
            let val = self.parse_value()?;
            self.module.elem_content[el_id as usize] = Some(val);
        }

        Ok(NodeId::Element(el_id))
    }

    pub fn parse_node(&mut self) -> Result<NodeId, String> {
        self.skip_whitespace();
        let c = self.peek().ok_or("EOF reached")?;

        match c {
            b'@' => self.parse_directive(),
            b'!' => {
                self.advance();
                let i = self.parse_ident();
                if i != "rhei" {
                    return Err("Expected rhei".into());
                }
                self.consume_if(b':');
                let expr = self.parse_rhei_expr()?;
                Ok(NodeId::Directive(
                    self.module.push_directive(Directive::RheiBlock(expr)),
                ))
            }
            b'A'..=b'Z' => self.parse_element(),
            b'"' | b'$' => {
                let val = self.parse_value()?;
                let el_id = self.module.push_element("#text".to_string());
                self.module.elem_content[el_id as usize] = Some(val);
                Ok(NodeId::Element(el_id))
            }
            _ => Err(format!("Unexpected token: {}", c as char)),
        }
    }

    pub fn parse_all(&mut self) -> Result<(), String> {
        self.skip_whitespace();
        while self.peek().is_some() {
            let node = self.parse_node()?;
            self.module.hierarchy.push(node);
            self.skip_whitespace();
        }
        Ok(())
    }
}
