// Validator for things inside of a typing.Literal[]
// which can be an int, a string, bytes or an Enum value (including `class Foo(str, Enum)` type enums)
use core::fmt::Debug;
use std::cell::OnceCell;
use std::sync::Arc;

use pyo3::prelude::*;
use pyo3::types::{PyDict, PyInt, PyList, PyString};
use pyo3::{PyTraverseError, PyVisit, intern};

use ahash::AHashMap;

use crate::build_tools::{py_schema_err, py_schema_error_type};
use crate::errors::{ErrorType, ValError, ValResult};
use crate::input::{Input, ValidationMatch};
use crate::py_gc::PyGcTraverse;
use crate::tools::SchemaDict;

use super::{BuildValidator, CombinedValidator, DefinitionsBuilder, ValidationState, Validator};

#[derive(Debug, Clone, Default)]
struct BoolLiteral {
    pub true_id: Option<usize>,
    pub false_id: Option<usize>,
}

#[derive(Debug, Clone)]
pub struct LiteralLookup<T: Debug> {
    // Specialized lookups for ints, bools and strings because they
    // (1) are easy to convert between Rust and Python
    // (2) hashing them in Rust is very fast
    // (3) are the most commonly used things in Literal[...]
    expected_bool: Option<BoolLiteral>,
    expected_int: Option<AHashMap<i64, usize>>,
    expected_str: Option<AHashMap<String, usize>>,
    // Catch all for hashable types like Enum and bytes (the latter only because it is seldom used)
    expected_py_dict: Option<Py<PyDict>>,
    // Catch all for unhashable types like list
    expected_py_values: Option<Vec<(Py<PyAny>, usize)>>,
    // Fallback for ints, bools, and strings to use Python hash and equality checks
    // which we can't mix with `expected_py_dict`, as there may be conflicts
    // for an example, see tests/test_validators/test_literal.py::test_mix_int_enum_with_int
    expected_py_primitives: Option<Py<PyDict>>,

    pub values: Vec<T>,
}

impl<T: Debug> LiteralLookup<T> {
    pub fn new<'py>(py: Python<'py>, expected: impl Iterator<Item = (Bound<'py, PyAny>, T)>) -> PyResult<Self> {
        let mut expected_bool = BoolLiteral::default();
        let mut expected_int = AHashMap::new();
        let mut expected_str: AHashMap<String, usize> = AHashMap::new();
        let expected_py_dict = PyDict::new(py);
        let mut expected_py_values = Vec::new();
        let expected_py_primitives = PyDict::new(py);
        let mut values = Vec::new();
        for (k, v) in expected {
            let id = values.len();
            values.push(v);

            if let Ok(bool_value) = k.validate_bool(true) {
                if bool_value.into_inner() {
                    expected_bool.true_id = Some(id);
                } else {
                    expected_bool.false_id = Some(id);
                }
                expected_py_primitives.set_item(&k, id)?;
            }
            if k.is_exact_instance_of::<PyInt>() {
                if let Ok(int_64) = k.extract::<i64>() {
                    expected_int.insert(int_64, id);
                    expected_py_primitives.set_item(&k, id)?;
                } else {
                    // cover the case of an int that's > i64::MAX etc.
                    expected_py_dict.set_item(k, id)?;
                }
            } else if let Ok(either_str) = k.exact_str() {
                let str = either_str
                    .as_cow()
                    .map_err(|_| py_schema_error_type!("error extracting str {k:?}"))?;
                expected_str.insert(str.to_string(), id);
                expected_py_primitives.set_item(&k, id)?;
            } else if expected_py_dict.set_item(&k, id).is_err() {
                expected_py_values.push((k.as_unbound().clone_ref(py), id));
            }
        }

        Ok(Self {
            expected_bool: (expected_bool.true_id.is_some() || expected_bool.false_id.is_some())
                .then_some(expected_bool),
            expected_int: (!expected_int.is_empty()).then_some(expected_int),
            expected_str: (!expected_str.is_empty()).then_some(expected_str),
            expected_py_dict: (!expected_py_dict.is_empty()).then_some(expected_py_dict.into()),
            expected_py_values: (!expected_py_values.is_empty()).then_some(expected_py_values),
            expected_py_primitives: (!expected_py_primitives.is_empty()).then_some(expected_py_primitives.into()),
            values,
        })
    }

    pub fn validate<'a, 'py, I: Input<'py> + ?Sized>(
        &self,
        py: Python<'py>,
        input: &'a I,
    ) -> ValResult<Option<(&'a I, &T)>> {
        if let Some(expected_bool) = &self.expected_bool
            && let Ok(bool_value) = input.validate_bool(true)
        {
            if bool_value.into_inner() {
                if let Some(true_value) = &expected_bool.true_id {
                    return Ok(Some((input, &self.values[*true_value])));
                }
            } else if let Some(false_value) = &expected_bool.false_id {
                return Ok(Some((input, &self.values[*false_value])));
            }
        }
        if let Some(expected_ints) = &self.expected_int
            && let Ok(either_int) = input.exact_int()
        {
            let int = either_int.into_i64(py)?;
            if let Some(id) = expected_ints.get(&int) {
                return Ok(Some((input, &self.values[*id])));
            }
        }
        if let Some(expected_strings) = &self.expected_str {
            let validation_result = if input.as_python().is_some() {
                input.exact_str()
            } else {
                // Strings coming from JSON are treated as "strict" but not "exact" for reasons
                // of parsing types like UUID; see the implementation of `validate_str` for Json
                // inputs for justification. We might change that eventually, but for now we need
                // to work around this when loading from JSON
                // V3 TODO: revisit making this "exact" for JSON inputs
                input.validate_str(true, false).map(ValidationMatch::into_inner)
            };

            if let Ok(either_str) = validation_result {
                let cow = either_str.as_cow()?;
                if let Some(id) = expected_strings.get(cow.as_ref()) {
                    return Ok(Some((input, &self.values[*id])));
                }
            }
        }
        // cache py_input if needed, since we might need it for multiple lookups
        let py_input = OnceCell::new();
        let get_py_input = || match py_input.get() {
            Some(py_input) => PyResult::<_>::Ok(py_input),
            None => {
                let _ = py_input.set(input.to_object(py)?);
                Ok(py_input.get().unwrap())
            }
        };

        if let Some(expected_py_dict) = &self.expected_py_dict &&
            // We don't use ? to unpack the result of `get_item` in the next line because unhashable
            // inputs will produce a TypeError, which in this case we just want to treat equivalently
            // to a failed lookup
             let Ok(Some(v)) = expected_py_dict.bind(py).get_item(get_py_input()?)
        {
            let id: usize = v.extract().unwrap();
            return Ok(Some((input, &self.values[id])));
        }
        if let Some(expected_py_values) = &self.expected_py_values {
            let py_input = get_py_input()?;
            for (k, id) in expected_py_values {
                if k.bind(py).eq(py_input).unwrap_or(false) {
                    return Ok(Some((input, &self.values[*id])));
                }
            }
        }

        // this one must be last to avoid conflicts with the other lookups, think of this
        // almost as a lax fallback
        if let Some(expected_py_primitives) = &self.expected_py_primitives
            // We don't use ? to unpack the result of `get_item` in the next line because unhashable
            // inputs will produce a TypeError, which in this case we just want to treat equivalently
            // to a failed lookup
            && let Ok(Some(v)) = expected_py_primitives.bind(py).get_item(get_py_input()?)
        {
            let id: usize = v.extract().unwrap();
            return Ok(Some((input, &self.values[id])));
        }
        Ok(None)
    }

    /// Used by int enums
    pub fn validate_int<'a, 'py, I: Input<'py> + ?Sized>(
        &self,
        py: Python<'py>,
        input: &'a I,
        strict: bool,
    ) -> ValResult<Option<&T>> {
        if let Some(expected_ints) = &self.expected_int
            && let Ok(either_int) = input.validate_int(strict)
            && let Some(id) = expected_ints.get(&either_int.into_inner().into_i64(py)?)
        {
            return Ok(Some(&self.values[*id]));
        }
        Ok(None)
    }

    /// Used by str enums
    pub fn validate_str<'a, 'py, I: Input<'py> + ?Sized>(&self, input: &'a I, strict: bool) -> ValResult<Option<&T>> {
        if let Some(expected_strings) = &self.expected_str
            && let Ok(either_str) = input.validate_str(strict, false)
            && let Some(id) = expected_strings.get(either_str.into_inner().as_cow()?.as_ref())
        {
            return Ok(Some(&self.values[*id]));
        }
        Ok(None)
    }

    /// Used by float enums
    pub fn validate_float<'a, 'py, I: Input<'py> + ?Sized>(
        &self,
        py: Python<'py>,
        input: &'a I,
        strict: bool,
    ) -> ValResult<Option<&T>> {
        if let Some(expected_py) = &self.expected_py_dict
            && let Ok(either_float) = input.validate_float(strict)
            && let Ok(Some(v)) = expected_py.bind(py).get_item(either_float.into_inner().as_f64())
        {
            let id: usize = v.extract().unwrap();
            return Ok(Some(&self.values[id]));
        }
        Ok(None)
    }
}

impl<T: PyGcTraverse + Debug> PyGcTraverse for LiteralLookup<T> {
    fn py_gc_traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
        self.expected_py_dict.py_gc_traverse(visit)?;
        self.values.py_gc_traverse(visit)?;
        Ok(())
    }
}

#[derive(Debug, Clone)]
pub struct LiteralValidator {
    lookup: LiteralLookup<Py<PyAny>>,
    expected_repr: String,
    name: String,
}

impl BuildValidator for LiteralValidator {
    const EXPECTED_TYPE: &'static str = "literal";

    fn build(
        schema: &Bound<'_, PyDict>,
        _config: Option<&Bound<'_, PyDict>>,
        _definitions: &mut DefinitionsBuilder<Arc<CombinedValidator>>,
    ) -> PyResult<Arc<CombinedValidator>> {
        let expected: Bound<PyList> = schema.get_as_req(intern!(schema.py(), "expected"))?;
        if expected.is_empty() {
            return py_schema_err!("`expected` should have length > 0");
        }
        let py = expected.py();
        let repr_args = expected.iter().map(|item| item.repr()).collect::<PyResult<Vec<_>>>()?;
        let expected_repr = expected_repr(&repr_args)?;
        let name = expected_name(&repr_args, Self::EXPECTED_TYPE)?;
        let lookup = LiteralLookup::new(py, expected.into_iter().map(|v| (v.clone(), v.into())))?;
        Ok(CombinedValidator::Literal(Self {
            lookup,
            expected_repr,
            name,
        })
        .into())
    }
}

impl_py_gc_traverse!(LiteralValidator { lookup });

impl Validator for LiteralValidator {
    fn validate<'py>(
        &self,
        py: Python<'py>,
        input: &(impl Input<'py> + ?Sized),
        _state: &mut ValidationState<'_, 'py>,
    ) -> ValResult<Py<PyAny>> {
        match self.lookup.validate(py, input)? {
            Some((_, v)) => Ok(v.clone()),
            None => Err(ValError::new(
                ErrorType::LiteralError {
                    expected: self.expected_repr.clone(),
                    context: None,
                },
                input,
            )),
        }
    }

    fn get_name(&self) -> &str {
        &self.name
    }
}

/// Formats strings according to (`base_name[arg1,arg2,...,argN]`)
pub fn expected_name(repr_args: &[Bound<'_, PyString>], base_name: &'static str) -> PyResult<String> {
    let mut args_len = 0;
    for arg in repr_args {
        args_len += arg.to_str()?.len();
    }

    let name_capacity = base_name.len()
        + 2 // for the brackets
        + args_len
        + repr_args.len().saturating_sub(1); // for commas

    let mut name = String::with_capacity(name_capacity);
    name.push_str(base_name);
    name.push('[');

    let mut first = true;
    for arg in repr_args {
        if first {
            first = false;
        } else {
            name.push(',');
        }
        let s = arg.to_str()?;
        name.push_str(s);
    }

    // close the name bracket
    name.push(']');

    // Should have used all the allocated capacity
    debug_assert_eq!(name_capacity, name.len());

    Ok(name)
}

/// Formats strings according to (`arg1, arg2, ... or argN]`)
pub fn expected_repr(repr_args: &[Bound<'_, PyString>]) -> PyResult<String> {
    let mut args_len = 0;
    for arg in repr_args {
        args_len += arg.to_str()?.len();
    }

    let repr_capacity = args_len
        + if repr_args.len() > 1 {
            // for the " or "
            4
        } else {
            0
        }
        + repr_args.len().saturating_sub(2) * 2; // for the commas

    let mut repr = String::with_capacity(repr_capacity);

    let (last_arg, repr_args) = repr_args.split_last().unwrap();
    let mut first = true;
    for arg in repr_args {
        if first {
            first = false;
        } else {
            repr.push_str(", ");
        }
        let s = arg.to_str()?;
        repr.push_str(s);
    }

    if !repr_args.is_empty() {
        repr.push_str(" or ");
    }
    repr.push_str(last_arg.to_str()?);

    // Should have used all the allocated capacity
    debug_assert_eq!(repr_capacity, repr.len());

    Ok(repr)
}
