use crate::ast::{Number, Value};
use pest::iterators::Pair;
use pest::Parser;
use pest_derive::Parser;
use std::collections::HashMap;
use thiserror::Error;

#[derive(Parser)]
#[grammar = "uson.pest"]
pub struct UsonParser;

#[derive(Error, Debug)]
pub enum ParseError {
    #[error("Parse error: {0}")]
    PestError(#[from] pest::error::Error<Rule>),
    
    #[error("Invalid number format: {0}")]
    NumberError(String),
    
    #[error("Invalid escape sequence: {0}")]
    EscapeError(String),
}

pub type ParseResult<T> = Result<T, ParseError>;

pub fn parse(input: &str) -> ParseResult<Vec<Value>> {
    let mut pairs = UsonParser::parse(Rule::uson_text, input)?;
    let uson_text = pairs.next().unwrap();
    
    let mut values = Vec::new();
    for pair in uson_text.into_inner() {
        if pair.as_rule() == Rule::EOI {
            break;
        }
        if pair.as_rule() == Rule::expr {
            let value_pair = pair.into_inner().next().unwrap();
            values.push(parse_value(value_pair)?);
        }
    }
    
    Ok(values)
}

fn parse_value(pair: Pair<Rule>) -> ParseResult<Value> {
    match pair.as_rule() {
        Rule::null_val => Ok(Value::Null),
        Rule::true_val => Ok(Value::Bool(true)),
        Rule::false_val => Ok(Value::Bool(false)),
        Rule::number => parse_number(pair),
        Rule::string => parse_string(pair),
        Rule::array => parse_array(pair),
        Rule::object => parse_object(pair),
        Rule::assign => parse_assign(pair),
        Rule::typed => parse_typed(pair),
        Rule::value => {
            let inner = pair.into_inner().next().unwrap();
            parse_value(inner)
        }
        _ => unreachable!("Unexpected rule: {:?}", pair.as_rule()),
    }
}

fn parse_number(pair: Pair<Rule>) -> ParseResult<Value> {
    let num_str = pair.as_str();
    
    if num_str.contains('.') || num_str.contains('e') || num_str.contains('E') {
        let f = num_str.parse::<f64>()
            .map_err(|e| ParseError::NumberError(e.to_string()))?;
        Ok(Value::Number(Number::Float(f)))
    } else {
        let i = num_str.parse::<i64>()
            .map_err(|e| ParseError::NumberError(e.to_string()))?;
        Ok(Value::Number(Number::Integer(i)))
    }
}

fn parse_string(pair: Pair<Rule>) -> ParseResult<Value> {
    let inner = pair.into_inner().next().unwrap();
    let s = match inner.as_rule() {
        Rule::double_quoted_string => {
            let s = inner.as_str();
            unescape_string(&s[1..s.len()-1])?
        }
        Rule::single_quoted_string => {
            let s = inner.as_str();
            unescape_string(&s[1..s.len()-1])?
        }
        Rule::unquoted_string => {
            let s = inner.as_str();
            // Check if it's a number
            if s.chars().all(|c| c.is_ascii_digit() || c == '.') {
                if s.contains('.') {
                    if let Ok(f) = s.parse::<f64>() {
                        return Ok(Value::Number(Number::Float(f)));
                    }
                } else if let Ok(i) = s.parse::<i64>() {
                    return Ok(Value::Number(Number::Integer(i)));
                }
            }
            s.to_string()
        }
        _ => unreachable!(),
    };
    
    Ok(Value::String(s))
}

fn unescape_string(s: &str) -> ParseResult<String> {
    let mut result = String::new();
    let mut chars = s.chars();
    
    while let Some(ch) = chars.next() {
        if ch == '\\' {
            match chars.next() {
                Some('n') => result.push('\n'),
                Some('r') => result.push('\r'),
                Some('t') => result.push('\t'),
                Some('b') => result.push('\u{0008}'),
                Some('f') => result.push('\u{000C}'),
                Some('\\') => result.push('\\'),
                Some('/') => result.push('/'),
                Some('#') => result.push('#'),
                Some('"') => result.push('"'),
                Some('\'') => result.push('\''),
                Some('u') => {
                    let hex: String = chars.by_ref().take(4).collect();
                    let code = u32::from_str_radix(&hex, 16)
                        .map_err(|_| ParseError::EscapeError(format!("Invalid unicode escape: \\u{}", hex)))?;
                    let ch = char::from_u32(code)
                        .ok_or_else(|| ParseError::EscapeError(format!("Invalid unicode codepoint: {}", code)))?;
                    result.push(ch);
                }
                Some(c) => result.push(c),
                None => return Err(ParseError::EscapeError("Incomplete escape sequence".to_string())),
            }
        } else {
            result.push(ch);
        }
    }
    
    Ok(result)
}

fn parse_array(pair: Pair<Rule>) -> ParseResult<Value> {
    let mut values = Vec::new();
    
    for inner in pair.into_inner() {
        if inner.as_rule() == Rule::value {
            values.push(parse_value(inner)?);
        }
    }
    
    Ok(Value::Array(values))
}

fn parse_object(pair: Pair<Rule>) -> ParseResult<Value> {
    let mut map = HashMap::new();
    
    for inner in pair.into_inner() {
        if inner.as_rule() == Rule::member {
            let mut member_parts = inner.into_inner();
            let key_pair = member_parts.next().unwrap();
            let key = if let Value::String(s) = parse_string(key_pair)? {
                s
            } else {
                unreachable!()
            };
            let value_pair = member_parts.next().unwrap();
            let value = parse_value(value_pair)?;
            map.insert(key, value);
        }
    }
    
    Ok(Value::Object(map))
}

fn parse_assign(pair: Pair<Rule>) -> ParseResult<Value> {
    let mut inner = pair.into_inner();
    let key_pair = inner.next().unwrap();
    let key = if let Value::String(s) = parse_string(key_pair)? {
        s
    } else {
        unreachable!()
    };
    let value_pair = inner.next().unwrap();
    let value = parse_value(value_pair)?;
    
    let mut map = HashMap::new();
    map.insert(key, value);
    Ok(Value::Object(map))
}

fn parse_typed(pair: Pair<Rule>) -> ParseResult<Value> {
    let mut inner = pair.into_inner();
    let type_name = inner.next().unwrap().as_str().to_string();
    let value_pair = inner.next().unwrap();
    let value = parse_value(value_pair)?;
    
    Ok(Value::Typed {
        type_name,
        value: Box::new(value),
    })
}

pub fn apply_builtin_types(value: Value) -> Value {
    match value {
        Value::Typed { type_name, value } => {
            let inner = apply_builtin_types(*value);
            match type_name.as_str() {
                "str" => match &inner {
                    Value::String(s) => Value::String(s.clone()),
                    Value::Number(Number::Integer(i)) => Value::String(i.to_string()),
                    Value::Number(Number::Float(f)) => Value::String(f.to_string()),
                    Value::Bool(b) => Value::String(b.to_string()),
                    Value::Null => Value::String("null".to_string()),
                    Value::Array(arr) => {
                        let strs: Vec<String> = arr.iter().map(|v| match v {
                            Value::String(s) => s.clone(),
                            Value::Number(Number::Integer(i)) => i.to_string(),
                            Value::Number(Number::Float(f)) => f.to_string(),
                            _ => "".to_string(),
                        }).collect();
                        Value::String(strs.join(","))
                    }
                    _ => inner,
                },
                "int" => match &inner {
                    Value::String(s) => s.parse::<i64>()
                        .map(|i| Value::Number(Number::Integer(i)))
                        .unwrap_or(inner),
                    Value::Number(_) => inner,
                    _ => inner,
                },
                "float" => match &inner {
                    Value::String(s) => s.parse::<f64>()
                        .map(|f| Value::Number(Number::Float(f)))
                        .unwrap_or(inner),
                    Value::Number(Number::Integer(i)) => Value::Number(Number::Float(*i as f64)),
                    Value::Number(_) => inner,
                    _ => inner,
                },
                "bool" => match &inner {
                    Value::String(s) => match s.as_str() {
                        "true" => Value::Bool(true),
                        "false" => Value::Bool(false),
                        _ => inner,
                    },
                    Value::Bool(_) => inner,
                    _ => inner,
                },
                "null" => Value::Null,
                "arr" => match &inner {
                    Value::String(s) => {
                        let parts: Vec<Value> = s.split(',')
                            .filter_map(|p| p.trim().parse::<i64>().ok())
                            .map(|i| Value::Number(Number::Integer(i)))
                            .collect();
                        Value::Array(parts)
                    }
                    Value::Array(_) => inner,
                    _ => inner,
                },
                "obj" => match &inner {
                    Value::String(s) => {
                        // Try to parse as JSON
                        if let Ok(parsed) = serde_json::from_str::<serde_json::Value>(s) {
                            json_to_value(parsed)
                        } else {
                            inner
                        }
                    }
                    _ => inner,
                },
                _ => inner, // Unknown type, return value as-is
            }
        }
        Value::Array(arr) => Value::Array(arr.into_iter().map(apply_builtin_types).collect()),
        Value::Object(obj) => Value::Object(
            obj.into_iter()
                .map(|(k, v)| (k, apply_builtin_types(v)))
                .collect()
        ),
        other => other,
    }
}

fn json_to_value(json: serde_json::Value) -> Value {
    match json {
        serde_json::Value::Null => Value::Null,
        serde_json::Value::Bool(b) => Value::Bool(b),
        serde_json::Value::Number(n) => {
            if let Some(i) = n.as_i64() {
                Value::Number(Number::Integer(i))
            } else if let Some(f) = n.as_f64() {
                Value::Number(Number::Float(f))
            } else {
                Value::Null
            }
        }
        serde_json::Value::String(s) => Value::String(s),
        serde_json::Value::Array(arr) => {
            Value::Array(arr.into_iter().map(json_to_value).collect())
        }
        serde_json::Value::Object(obj) => {
            Value::Object(obj.into_iter().map(|(k, v)| (k, json_to_value(v))).collect())
        }
    }
}