use crate::ast::*;
use std::collections::HashMap;

pub struct StyleParser<'a, 'm> {
    src: &'a [u8],
    pos: usize,
    pub module: &'m mut ModuleSoA<'a>,
    variables: HashMap<String, Value<'a>>,
    mixins: HashMap<String, Vec<(&'a str, Value<'a>)>>,
}

impl<'a, 'm> StyleParser<'a, 'm> {
    pub fn new(src: &'a str, module: &'m mut ModuleSoA<'a>) -> Self {
        Self {
            src: src.as_bytes(),
            pos: 0,
            module,
            variables: HashMap::new(),
            mixins: HashMap::new(),
        }
    }

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

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

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

    #[inline]
    fn skip_whitespace(&mut self) {
        let len = self.src.len();
        while self.pos < len {
            let c = self.src[self.pos];
            if c.is_ascii_whitespace() {
                self.pos += 1;
            } else if c == b'/' && self.pos + 1 < len && self.src[self.pos + 1] == b'/' {
                self.pos += 2;
                while self.pos < len && self.src[self.pos] != b'\n' {
                    self.pos += 1;
                }
            } else {
                break;
            }
        }
    }

    #[inline]
    fn parse_ident(&mut self) -> &'a str {
        let start = self.pos;
        while self.pos < self.src.len() {
            let c = self.src[self.pos];
            if c.is_ascii_alphanumeric() || c == b'-' || c == b'_' {
                self.pos += 1;
            } else {
                break;
            }
        }
        unsafe { std::str::from_utf8_unchecked(&self.src[start..self.pos]) }
    }

    #[inline]
    fn parse_selector(&mut self) -> &'a str {
        let start = self.pos;
        while self.pos < self.src.len() && self.src[self.pos] != b'{' {
            self.pos += 1;
        }
        unsafe { std::str::from_utf8_unchecked(&self.src[start..self.pos]) }.trim()
    }

    fn parse_single_value(&self, s: &'a str) -> Value<'a> {
        let s = s.trim_end_matches(',');
       
        if let Some(_stripped) = s.strip_prefix('#') {
            return Value::Color(s);
        }
        
        if let Some(var_name) = s.strip_prefix('$') {
            if let Some(val) = self.variables.get(var_name) {
                return val.clone();
            }
            return Value::Variable(var_name);
        }

        if s.contains('(') && s.ends_with(')') {
            let parts: Vec<&str> = s.splitn(2, '(').collect();
            let name = parts[0].trim();
            let args_str = parts[1].trim_end_matches(')');
            let args = args_str.split(',')
                .filter(|a| !a.trim().is_empty())
                .map(|a| self.parse_single_value(a.trim()))
                .collect();
            return Value::Call(name, args);
        }

        let mut num_end = 0;
        for (i, c) in s.char_indices() {
            if c.is_ascii_digit() || c == '.' || c == '-' {
                num_end = i + c.len_utf8();
            } else { break; }
        }
        
        if num_end > 0 && num_end < s.len() {
            if let Ok(num) = s[..num_end].parse::<f64>() {
                return Value::Unit(num, &s[num_end..]);
            }
        }

        if let Ok(i) = s.parse::<i64>() { return Value::Int(i); }
        if let Ok(f) = s.parse::<f64>() { return Value::Float(f); }
        
        match s {
            "true" => Value::Bool(true),
            "false" => Value::Bool(false),
            "null" => Value::Null,
            _ => Value::Ident(s),
        }
    }

    fn parse_value_line(&mut self) -> Value<'a> {
        let start = self.pos;
        let mut parens = 0;
        let len = self.src.len();
        
        while self.pos < len {
            let c = self.src[self.pos];
            if c == b'(' { parens += 1; }
            else if c == b')' { parens -= 1; }
            else if parens == 0 && matches!(c, b'\n' | b'\r' | b';' | b'}') {
                break;
            }
            self.pos += 1;
        }

        let line = unsafe { std::str::from_utf8_unchecked(&self.src[start..self.pos]) }.trim();
        if self.peek() == Some(b';') { self.advance(); }

        let mut tokens = Vec::new();
        let mut token_start = None;
        let mut p = 0;
        
        let bytes = line.as_bytes();
        for (i, &c) in bytes.iter().enumerate() {
            match c {
                b'(' => {
                    p += 1;
                    if token_start.is_none() { token_start = Some(i); }
                }
                b')' => {
                    if p > 0 { p -= 1; }
                    if token_start.is_none() { token_start = Some(i); }
                }
                b' ' | b'\t' if p == 0 => {
                    if let Some(s_idx) = token_start {
                        tokens.push(&line[s_idx..i]);
                        token_start = None;
                    }
                }
                _ => {
                    if token_start.is_none() { token_start = Some(i); }
                }
            }
        }
        if let Some(s_idx) = token_start {
            tokens.push(&line[s_idx..]);
        }

        if tokens.len() == 1 {
            self.parse_single_value(tokens[0])
        } else {
            Value::Array(tokens.into_iter().map(|t| self.parse_single_value(t)).collect())
        }
    }

    fn parse_properties(&mut self, target: &mut Vec<(&'a str, Value<'a>)>) -> Result<(), String> {
        self.consume_if(b'{');
        loop {
            self.skip_whitespace();
            if self.consume_if(b'}') { break; }

            if self.peek() == Some(b'@') {
                self.advance();
                let ident = self.parse_ident();
                if ident == "use" {
                    self.skip_whitespace();
                    let mixin_name = self.parse_ident();
                    if let Some(props) = self.mixins.get(mixin_name) {
                        target.extend(props.clone());
                    } else {
                        return Err(format!("Unknown mixin: {}", mixin_name));
                    }
                }
                continue;
            }

            let key = self.parse_ident();
            if key.is_empty() { return Err("Expected property key".into()); }
            
            self.skip_whitespace();
            if !self.consume_if(b':') { return Err("Expected ':' after property".into()); }
            self.skip_whitespace();

            let val = self.parse_value_line();
            target.push((key, val));
        }
        Ok(())
    }

    fn parse_top_level(&mut self) -> Result<Option<NodeId>, String> {
        match self.peek() {
            Some(b'$') => {
                self.advance();
                let name = self.parse_ident().to_string();
                self.skip_whitespace();
                self.consume_if(b'=');
                self.skip_whitespace();
                let val = self.parse_value_line();
                self.variables.insert(name, val);
                Ok(None)
            }
            Some(b'@') => {
                self.advance();
                let name = self.parse_ident();
                match name {
                    "mixin" => {
                        self.skip_whitespace();
                        let mixin_name = self.parse_ident().to_string();
                        self.skip_whitespace();
                        let mut props = Vec::new();
                        self.parse_properties(&mut props)?;
                        self.mixins.insert(mixin_name, props);
                        Ok(None)
                    }
                    "anim" => {
                        self.skip_whitespace();
                        let anim_name = self.parse_ident();
                        self.skip_whitespace();
                        self.consume_if(b'{');
                        
                        let mut frames = Vec::new();
                        loop {
                            self.skip_whitespace();
                            if self.consume_if(b'}') { break; }
                            
                            let step = self.parse_selector().to_string();
                            let mut props = Vec::new();
                            self.parse_properties(&mut props)?;
                            
                            let start_idx = self.module.prop_keys.len() as u32;
                            for (k, v) in props {
                                self.module.prop_keys.push(k);
                                self.module.prop_values.push(v);
                            }
                            let len = self.module.prop_keys.len() as u32 - start_idx;
                            frames.push((step, (start_idx, len)));
                        }
                        
                        let id = self.module.push_directive(Directive::StyleAnim { name: anim_name, frames });
                        Ok(Some(NodeId::Directive(id)))
                    }
                    "if" => {
                        self.skip_whitespace();
                        let start = self.pos;
                        while self.pos < self.src.len() && self.src[self.pos] != b'{' { self.pos += 1; }
                        let cond_str = unsafe { std::str::from_utf8_unchecked(&self.src[start..self.pos]) }.trim();
                        
                        self.consume_if(b'{');
                        let start_idx = self.module.hierarchy.len() as u32;
                        loop {
                            self.skip_whitespace();
                            if self.consume_if(b'}') { break; }
                            if let Some(node) = self.parse_top_level()? {
                                self.module.hierarchy.push(node);
                            }
                        }
                        let len = self.module.hierarchy.len() as u32 - start_idx;
                        
                        let condition = Value::Rhei(cond_str);
                        let id = self.module.push_directive(Directive::If { 
                            condition, 
                            child_span: (start_idx, len), 
                            else_span: None 
                        });
                        Ok(Some(NodeId::Directive(id)))
                    }
                    _ => Err(format!("Unknown directive: @{}", name)),
                }
            }
            _ => {
                let selector = self.parse_selector();
                if selector.is_empty() { return Err("Expected selector".into()); }
                
                let mut props = Vec::new();
                self.parse_properties(&mut props)?;

                let start_idx = self.module.prop_keys.len() as u32;
                for (k, v) in props {
                    self.module.prop_keys.push(k);
                    self.module.prop_values.push(v);
                }
                let len = self.module.prop_keys.len() as u32 - start_idx;

                let id = self.module.push_directive(Directive::StyleRule { selector, prop_span: (start_idx, len) });
                Ok(Some(NodeId::Directive(id)))
            }
        }
    }

    pub fn parse_all(&mut self) -> Result<(), String> {
        self.skip_whitespace();
        while self.pos < self.src.len() {
            if let Some(node) = self.parse_top_level()? {
                self.module.hierarchy.push(node);
            }
            self.skip_whitespace();
        }
        Ok(())
    }
}