use std::{
    fmt::{self, Debug},
    sync::{Arc, OnceLock},
};

use itertools::Itertools;

use crate::lexer::token::{Keyword, TokenTag};

#[derive(Clone)]
pub enum Expected<K: Keyword> {
    Name(&'static str),
    Token(TokenTag<K>),
    Sequence(Vec<Expected<K>>),
    Alternatives(Vec<Expected<K>>),
    Optional(Box<Expected<K>>),

    /// Lazily computed expected. The value is memoized (computed at most once).
    ///
    /// NOTE: This can still represent cyclic structures; formatting/evaluation must guard
    /// against runaway recursion.
    Defered(DeferedExpected<K>),
}

#[derive(Clone)]
pub struct DeferedExpected<K: Keyword> {
    memo: Arc<OnceLock<Expected<K>>>,
    f: Arc<dyn Fn() -> Expected<K>>,
}

impl<K: Keyword> DeferedExpected<K> {
    pub fn new<F: Fn() -> Expected<K> + 'static>(f: F) -> Self {
        Self {
            memo: Arc::new(OnceLock::new()),
            // We require Send+Sync because OnceLock is Sync and this is generally safer.
            // If your Expected is truly single-threaded, you can drop these bounds.
            f: Arc::new(f),
        }
    }

    /// Get the memoized value, computing it once if needed.
    #[inline]
    pub fn get(&self) -> &Expected<K> {
        self.memo.get_or_init(|| (self.f)())
    }
}

// ---- formatting recursion guard ----
//
// We guard expansions of deferred values per-thread.
// This avoids stack overflows when Expected graphs are cyclic.
//
// The guard is intentionally small (3) to behave like Rust diagnostics
// (show a bit, then truncate).
const FORMAT_DEFERRED_EXPANSION_LIMIT: usize = 3;

thread_local! {
    static DEFERRED_FMT_DEPTH: std::cell::Cell<usize> = std::cell::Cell::new(0);
}

struct DeferredFmtGuard;

impl DeferredFmtGuard {
    fn enter() -> Option<Self> {
        let ok = DEFERRED_FMT_DEPTH.with(|d| {
            let cur = d.get();
            if cur >= FORMAT_DEFERRED_EXPANSION_LIMIT {
                false
            } else {
                d.set(cur + 1);
                true
            }
        });

        ok.then_some(Self)
    }
}

impl Drop for DeferredFmtGuard {
    fn drop(&mut self) {
        DEFERRED_FMT_DEPTH.with(|d| d.set(d.get().saturating_sub(1)));
    }
}

impl<K: Keyword> Debug for Expected<K> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Expected::Name(name) => write!(f, "{}", name),
            Expected::Token(token) => write!(f, "{}", token),
            Expected::Sequence(items) => write!(f, "{:?}", items),
            Expected::Alternatives(items) => write!(f, "{:?}", items),
            Expected::Optional(inner) => write!(f, "{:?}", inner),
            Expected::Defered(defered) => {
                // Debug printing should also be protected.
                if let Some(_g) = DeferredFmtGuard::enter() {
                    write!(f, "{:?}", defered.get())
                } else {
                    write!(f, "<defered ...>")
                }
            }
        }
    }
}

impl<K: Keyword> Expected<K> {
    pub fn format(&self, found: TokenTag<K>, first: bool) -> String {
        let one_of = matches!(self, Expected::Alternatives(items) if items.len() > 1);
        if one_of {
            format!(
                "expected one of {}, found {}",
                self.raw_format(first),
                found
            )
        } else {
            format!("expected {}, found {}", self.raw_format(first), found)
        }
    }

    pub fn first(&self) -> Expected<K> {
        match self {
            Expected::Optional(inner) => inner.first(),
            Expected::Name(_) => self.clone(),
            Expected::Sequence(items) => items
                .first()
                .cloned()
                .unwrap_or_else(|| Expected::Sequence(Vec::new())),
            Expected::Alternatives(items) => {
                Expected::Alternatives(items.iter().map(|alt| alt.first()).collect())
            }
            Expected::Token(_) => self.clone(),
            Expected::Defered(inner) => inner.get().first(),
        }
    }

    pub fn last(&self) -> Expected<K> {
        match self {
            Expected::Optional(inner) => inner.last(),
            Expected::Name(_) => self.clone(),
            Expected::Sequence(items) => items
                .last()
                .cloned()
                .unwrap_or_else(|| Expected::Sequence(Vec::new())),
            Expected::Alternatives(items) => items
                .last()
                .cloned()
                .unwrap_or_else(|| Expected::Alternatives(Vec::new())),
            Expected::Token(_) => self.clone(),
            Expected::Defered(inner) => inner.get().last(),
        }
    }

    pub fn raw_format(&self, first: bool) -> String {
        match self {
            Expected::Optional(inner) => format!("{}?", inner.raw_format(first)),
            Expected::Name(name) => name.to_string(),
            Expected::Token(expected) => expected.to_string(),

            Expected::Defered(inner) => {
                // Guard: expand only a few levels, then truncate.
                if let Some(_g) = DeferredFmtGuard::enter() {
                    inner.get().raw_format(first)
                } else {
                    "<defered ...>".to_string()
                }
            }

            Expected::Sequence(items) => {
                if first {
                    items
                        .first()
                        .map(|i| i.raw_format(true))
                        .unwrap_or_default()
                } else if items.len() <= 1 {
                    format!(
                        "{}",
                        items
                            .iter()
                            .format_with("", |i, f| f(&format_args!("{}", i.raw_format(first))))
                    )
                } else {
                    let parts: Vec<_> = items.iter().map(|i| i.raw_format(first)).collect();
                    format!("({})", parts.join(" "))
                }
            }

            Expected::Alternatives(items) => {
                let mut parts = items.iter().map(|i| i.raw_format(first));
                match items.len() {
                    0 => String::new(),
                    1 => parts.next().unwrap(),
                    _ => {
                        let all: Vec<_> = parts.collect();
                        let (last, rest) = all.split_last().unwrap();
                        format!("{} or {}", rest.join(", "), last)
                    }
                }
            }
        }
    }

    pub fn optional(self) -> Self {
        Expected::Optional(Box::new(self))
    }

    pub fn tok(kind: TokenTag<K>) -> Self {
        Expected::Token(kind)
    }

    pub fn seq(items: impl IntoIterator<Item = Expected<K>>) -> Self {
        Expected::Sequence(items.into_iter().collect())
    }

    pub fn alt(items: impl IntoIterator<Item = Expected<K>>) -> Self {
        Expected::Alternatives(items.into_iter().collect())
    }

    pub fn defered<F: Fn() -> Expected<K> + 'static>(f: F) -> Self {
        Expected::Defered(DeferedExpected::new(f))
    }

    /// Flatten nested sequences into a single sequence
    pub fn then(self, other: Expected<K>) -> Self {
        match (self, other) {
            (Expected::Sequence(mut a), Expected::Sequence(b)) => {
                a.extend(b);
                Expected::Sequence(a)
            }
            (Expected::Sequence(mut a), b) => {
                a.push(b);
                Expected::Sequence(a)
            }
            (a, Expected::Sequence(mut b)) => {
                b.insert(0, a);
                Expected::Sequence(b)
            }
            (a, b) => Expected::Sequence(vec![a, b]),
        }
    }

    /// Flatten nested alternatives into a single alternative
    pub fn or(self, other: Expected<K>) -> Self {
        match (self, other) {
            (Expected::Alternatives(mut a), Expected::Alternatives(b)) => {
                a.extend(b);
                Expected::Alternatives(a)
            }
            (Expected::Alternatives(mut a), b) => {
                a.push(b);
                Expected::Alternatives(a)
            }
            (a, Expected::Alternatives(mut b)) => {
                b.insert(0, a);
                Expected::Alternatives(b)
            }
            (a, b) => Expected::Alternatives(vec![a, b]),
        }
    }
}
