use std::borrow::Cow;

use memchr::{memchr, memchr2};
use pyo3::{
    IntoPyObjectExt,
    prelude::*,
    types::{PyDict, PyFrozenSet, PyList, PySet, PyTuple},
};
use rustc_hash::FxHashMap;
use saphyr_parser::{Event, Parser, ScalarStyle, ScanError, Tag};

use crate::load::{
    arena::{Arena, NodeId},
    parse_datetime::parse_py_datetime,
    scalar::{is_bool, is_datetime, is_float, is_inf_nan, is_int, is_null, parse_float, parse_int},
    value::Value,
};

#[derive(Debug)]
enum Frame {
    Seq {
        anchor: usize,
        items: Vec<NodeId>,
    },
    Map {
        anchor: usize,
        items: Vec<(NodeId, NodeId)>,
        pending_key: Option<NodeId>,
    },
}

#[derive(Debug)]
pub enum BuildError {
    Scan(ScanError),
    Decode(String),
}

impl From<ScanError> for BuildError {
    fn from(err: ScanError) -> Self {
        BuildError::Scan(err)
    }
}

fn resolve_scalar<'a>(
    arena: &mut Arena<'a>,
    value: Cow<'a, str>,
    style: ScalarStyle,
    tag: Option<&Tag>,
) -> Result<NodeId, String> {
    if let Some(tag) = tag {
        if tag.is_yaml_core_schema() {
            let v = match tag.suffix.as_str() {
                "int" => parse_int(value.as_ref())
                    .ok_or_else(|| format!("Invalid value '{}' for '!!int' tag", value.as_ref()))?,
                "float" => parse_float(value.as_ref())
                    .map(Value::Float)
                    .ok_or_else(|| {
                        format!("Invalid value '{}' for '!!float' tag", value.as_ref())
                    })?,
                "bool" => is_bool(value.as_ref()).map(Value::Boolean).ok_or_else(|| {
                    format!("Invalid value '{}' for '!!bool' tag", value.as_ref())
                })?,
                "null" => {
                    let str = value.as_ref();
                    if str.is_empty() || is_null(str) {
                        Value::Null
                    } else {
                        return Err(format!("Invalid value '{str}' for '!!null' tag"));
                    }
                }
                "binary" => Value::String(value),
                "str" => Value::StringExplicit(value),
                _ => return Err(format!("Invalid tag: '!!{}'", tag.suffix)),
            };
            return Ok(arena.push(v));
        }

        return Ok(arena.push(Value::String(value)));
    }

    if style == ScalarStyle::Plain {
        let str = value.as_ref();

        if str.is_empty() || is_null(str) {
            return Ok(arena.push(Value::Null));
        }

        if let Some(bool) = is_bool(str) {
            return Ok(arena.push(Value::Boolean(bool)));
        }

        let bytes = str.as_bytes();

        if (is_inf_nan(bytes).is_some()
            || memchr(b'.', bytes).is_some()
            || memchr2(b'e', b'E', bytes).is_some())
            && is_float(bytes)
            && let Some(float) = parse_float(str)
        {
            return Ok(arena.push(Value::Float(float)));
        }

        if is_int(bytes)
            && let Some(int) = parse_int(str)
        {
            return Ok(arena.push(int));
        }
    }

    Ok(arena.push(Value::String(value)))
}

pub(crate) fn build_from_events(input: &'_ str) -> Result<(Arena<'_>, Vec<NodeId>), BuildError> {
    let parser = Parser::new_from_str(input);

    let mut arena = Arena::with_capacity((input.len() / 8).max(64));

    let mut stack: Vec<Frame> = Vec::new();
    let mut docs: Vec<NodeId> = Vec::new();
    let mut anchors: FxHashMap<usize, NodeId> = FxHashMap::default();
    let mut current_root: Option<NodeId> = None;

    for event_res in parser {
        let (event, _) = event_res?;

        match event {
            Event::StreamStart | Event::StreamEnd | Event::Nothing => {}
            Event::DocumentStart(_) => {
                current_root = None;
                stack.clear();
            }
            Event::DocumentEnd => {
                let root = current_root
                    .take()
                    .unwrap_or_else(|| arena.push(Value::Null));
                docs.push(root);
            }
            Event::Alias(id) => {
                let node = anchors
                    .get(&id)
                    .copied()
                    .unwrap_or_else(|| arena.push(Value::Null));

                push_value(node, &mut stack, &mut current_root);
            }
            Event::Scalar(val, style, anchor_id, tag) => {
                let node = resolve_scalar(&mut arena, val, style, tag.as_deref())
                    .map_err(BuildError::Decode)?;

                if anchor_id != 0 {
                    anchors.insert(anchor_id, node);
                }

                push_value(node, &mut stack, &mut current_root);
            }
            Event::SequenceStart(anchor_id, _) => {
                stack.push(Frame::Seq {
                    anchor: anchor_id,
                    items: Vec::new(),
                });
            }
            Event::SequenceEnd => {
                if let Some(Frame::Seq { anchor, items }) = stack.pop() {
                    let node = arena.push(Value::Seq(items));

                    if anchor != 0 {
                        anchors.insert(anchor, node);
                    }

                    push_value(node, &mut stack, &mut current_root);
                }
            }
            Event::MappingStart(anchor_id, _) => {
                stack.push(Frame::Map {
                    anchor: anchor_id,
                    items: Vec::new(),
                    pending_key: None,
                });
            }
            Event::MappingEnd => {
                if let Some(Frame::Map { anchor, items, .. }) = stack.pop() {
                    let node = arena.push(Value::Map(items));

                    if anchor != 0 {
                        anchors.insert(anchor, node);
                    }

                    push_value(node, &mut stack, &mut current_root);
                }
            }
        }
    }

    Ok((arena, docs))
}

#[inline]
fn push_value(value: NodeId, stack: &mut [Frame], root: &mut Option<NodeId>) {
    if let Some(top) = stack.last_mut() {
        match top {
            Frame::Seq { items, .. } => items.push(value),
            Frame::Map {
                items, pending_key, ..
            } => {
                if let Some(key) = pending_key.take() {
                    items.push((key, value));
                } else {
                    *pending_key = Some(value);
                }
            }
        }
    } else {
        *root = Some(value);
    }
}

fn value_to_py<'py>(
    py: Python<'py>,
    arena: &Arena<'_>,
    id: NodeId,
    parse_datetime: bool,
) -> PyResult<Bound<'py, PyAny>> {
    match arena.get(id) {
        Value::Null => Ok(py.None().into_bound(py)),
        Value::Boolean(bool) => bool.into_bound_py_any(py),
        Value::IntegerI64(int_64) => int_64.into_bound_py_any(py),
        Value::IntegerBig(big_int) => big_int.into_bound_py_any(py),
        Value::Float(float) => float.into_bound_py_any(py),
        Value::StringExplicit(string_exp) => string_exp.into_bound_py_any(py),
        Value::String(string) => {
            let str = string.as_ref();
            if parse_datetime
                && is_datetime(str.as_bytes())
                && let Ok(Some(dt)) = parse_py_datetime(py, str)
            {
                return Ok(dt);
            }
            str.into_bound_py_any(py)
        }
        Value::Seq(items) => {
            let py_list = PyList::empty(py);
            for &child in items {
                py_list.append(value_to_py(py, arena, child, parse_datetime)?)?;
            }
            Ok(py_list.into_any())
        }
        Value::Map(pairs) => {
            let mut all_nulls = true;
            let mut has_null_key = false;

            for (k, v) in pairs {
                if matches!(arena.get(*k), Value::Null) {
                    has_null_key = true;
                }
                if !matches!(arena.get(*v), Value::Null) {
                    all_nulls = false;
                }
            }

            if all_nulls && !has_null_key && pairs.len() > 1 {
                let py_set = PySet::empty(py)?;
                for (k, _) in pairs {
                    py_set.add(value_to_hashable(py, arena, *k, parse_datetime)?)?;
                }
                Ok(py_set.into_any())
            } else {
                let py_dict = PyDict::new(py);
                for (k, v) in pairs {
                    py_dict.set_item(
                        value_to_hashable(py, arena, *k, parse_datetime)?,
                        value_to_py(py, arena, *v, parse_datetime)?,
                    )?;
                }
                Ok(py_dict.into_any())
            }
        }
    }
}

fn value_to_hashable<'py>(
    py: Python<'py>,
    arena: &Arena<'_>,
    id: NodeId,
    parse_datetime: bool,
) -> PyResult<Bound<'py, PyAny>> {
    match arena.get(id) {
        Value::Seq(items) => {
            let mut vec = Vec::with_capacity(items.len());
            for &child in items {
                vec.push(value_to_hashable(py, arena, child, parse_datetime)?);
            }
            PyTuple::new(py, &vec)?.into_bound_py_any(py)
        }
        Value::Map(pairs) => {
            let py_list = PyList::empty(py);
            for (k, v) in pairs {
                let py_tuple = PyTuple::new(
                    py,
                    &[
                        value_to_hashable(py, arena, *k, parse_datetime)?,
                        value_to_py(py, arena, *v, parse_datetime)?,
                    ],
                )?;
                py_list.append(py_tuple)?;
            }
            PyFrozenSet::new(py, py_list)?.into_bound_py_any(py)
        }
        _ => value_to_py(py, arena, id, parse_datetime),
    }
}

pub(crate) fn to_python<'py>(
    py: Python<'py>,
    arena: &Arena<'_>,
    docs: &[NodeId],
    parse_datetime: bool,
) -> PyResult<Bound<'py, PyAny>> {
    match docs.len() {
        0 => Ok(py.None().into_bound(py)),
        1 => value_to_py(py, arena, docs[0], parse_datetime),
        _ => {
            let py_list = PyList::empty(py);
            for &doc in docs {
                py_list.append(value_to_py(py, arena, doc, parse_datetime)?)?;
            }
            Ok(py_list.into_any())
        }
    }
}
