// `;` are encoded in spans of statement nodes as very last token
// Need to do a proper check to figure out what should and shouldn't be a node
use crate::arena::ArenaVec;

use super::lexer::TokenPosition;

use crate::arena::{Arena, ArenaNode, ArenaString};

pub mod callables;
pub mod expressions;
pub mod operators;
pub mod types;

pub use callables::*;
pub use expressions::*;
pub use operators::*;
pub use types::*;

// Get rid of identifier field
#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq)]
pub struct IdentifierToken(pub TokenPosition);

#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq)]
pub struct OperatorToken(pub TokenPosition);

#[derive(Debug, Hash, PartialEq, Eq)]
pub struct QualifiedIdentifier<'arena> {
    pub head: IdentifierToken,
    pub tail: Option<ArenaVec<'arena, IdentifierToken>>, // None => single segment
}
pub type QualifiedIdentifierRef<'arena> = ArenaNode<'arena, QualifiedIdentifier<'arena>>;

// All inclusive!
#[derive(Debug, Hash, PartialEq, Eq, PartialOrd, Ord, Clone, Copy)]
pub struct AstSpan {
    pub token_from: TokenPosition,
    pub token_to: TokenPosition,
}

impl AstSpan {
    // -------- existing coord-based API (unchanged externally) --------

    #[must_use]
    pub const fn merge(left_span: &Self, right_span: &Self) -> Self {
        Self {
            // assumes both were constructed in the same style; good enough for the refactor
            token_from: left_span.token_from,
            token_to: right_span.token_to,
        }
    }

    // -------- NEW: 4 constructors based on TokenIndex --------

    /// Single-token span from an index (coords are dummy for now).
    #[inline]
    #[must_use]
    pub const fn new(single_index: TokenPosition) -> Self {
        Self {
            token_from: single_index,
            token_to: single_index,
        }
    }

    /// Span from two indices (coords are dummy for now).
    #[inline]
    #[must_use]
    pub const fn range(from: TokenPosition, to: TokenPosition) -> Self {
        Self {
            token_from: from,
            token_to: to,
        }
    }

    /// Immutable extension by index (keeps coords as-is).
    #[inline]
    #[must_use]
    pub fn extended(&self, right_most_index: TokenPosition) -> Self {
        Self {
            token_from: self.token_from,
            token_to: std::cmp::max(self.token_to, right_most_index),
        }
    }

    /// In-place extension by index (coords unchanged).
    #[inline]
    pub fn extend_to(&mut self, right_most_index: TokenPosition) {
        if right_most_index > self.token_to {
            self.token_to = right_most_index;
        }
    }
}

impl<'arena> QualifiedIdentifier<'arena> {
    #[inline]
    #[must_use]
    pub const fn is_single(&self) -> bool {
        self.tail.is_none()
    }

    #[inline]
    #[allow(clippy::len_without_is_empty)] // Suppress useless suggestion for `is_empty()`
    #[must_use]
    pub fn len(&self) -> usize {
        1 + self.tail.as_ref().map_or(0, |v| v.len())
    }

    #[inline]
    #[must_use]
    pub const fn head(&self) -> IdentifierToken {
        self.head
    }

    /// Iterates all identifier segments in order without allocating.
    pub fn iter(&self) -> impl Iterator<Item = IdentifierToken> + '_ {
        core::iter::once(self.head).chain(self.tail.iter().flat_map(|v| v.iter().copied()))
    }

    /// Cheap constructor from a single identifier. No Vec allocated.
    pub fn from_ident(arena: &'arena Arena, id: IdentifierToken) -> QualifiedIdentifierRef<'arena> {
        let span = AstSpan::new(id.0);
        ArenaNode::new_in(
            Self {
                head: id,
                tail: None,
            },
            span,
            arena,
        )
    }
    /// Cheap constructor from a single identifier. No Vec allocated.
    pub fn from_position(
        arena: &'arena Arena,
        position: TokenPosition,
    ) -> QualifiedIdentifierRef<'arena> {
        let span = AstSpan::new(position);
        ArenaNode::new_in(
            Self {
                head: IdentifierToken(position),
                tail: None,
            },
            span,
            arena,
        )
    }
}

#[derive(Debug, PartialEq)]
pub enum Statement<'src, 'arena> {
    // For the cases where user just used too many semi-colons `;;;;`
    Empty,
    Expression(ExpressionRef<'src, 'arena>),
    // Just declarations without assignment:
    // `local int i, j, k`
    LocalVariableDeclaration {
        type_spec: TypeSpecifierRef<'src, 'arena>,
        declarators: ArenaVec<'arena, VariableDeclaratorRef<'src, 'arena>>, // CHANGED
    },
    Label(ArenaString<'arena>),
    /// Nested function definitions inside blocks or states.
    Function(CallableDefinitionRef<'src, 'arena>),
    // For injecting in place of parts that couldn't be parsed
    Error,
}
pub type StatementRef<'src, 'arena> = ArenaNode<'arena, Statement<'src, 'arena>>;

#[derive(Debug)]
pub enum DeclarationLiteral<'src, 'arena> {
    None,
    Bool(bool),
    Integer(i128),
    Float(f64),
    String(ArenaString<'arena>),
    Identifier(&'src str),
    TaggedName {
        tag: IdentifierToken,
        quoted: ArenaString<'arena>,
    }, // NEW
}

#[derive(Debug)]
pub struct DeclarationLiteralRef<'src, 'arena> {
    pub literal: DeclarationLiteral<'src, 'arena>,
    pub position: TokenPosition,
}

impl IdentifierToken {
    #[must_use]
    pub const fn span(self) -> AstSpan {
        AstSpan::new(self.0)
    }
}

pub enum ClassModifier<'arena> {
    Final,
    Native,
    Abstract,
    Transient,
    Public,
    Protected,
    Private,
    Static,
    Config(Option<IdentifierToken>),
    NativeReplication,
    ExportStructs,
    SafeReplace,

    Const,
    Deprecated,
    NoExport,
    Export,

    Localized,
    Placeable,
    NotPlaceable,
    Instanced,
    EditConst,
    EditInline,
    EditInlineNew,
    NotEditInlineNew,
    CollapseCategories,
    DontCollapseCategories,
    HideCategories(ArenaVec<'arena, IdentifierToken>),
    ShowCategories(ArenaVec<'arena, IdentifierToken>),
    Within(IdentifierToken),
    DependsOn(IdentifierToken),
    GlobalConfig,
    PerObjectConfig,
    DynamicRecompile,
    HideDropdown,
    ParseConfig,
    CacheExempt,
}

pub type ClassModifierRef<'arena> = ArenaNode<'arena, ClassModifier<'arena>>;

pub struct ClassDeclaration<'arena> {
    pub name: IdentifierToken,
    pub parent: Option<IdentifierToken>,
    pub modifiers: Vec<ClassModifierRef<'arena>>,
}

// --- in ast.rs ---

#[derive(Debug)]
pub struct ClassVarDecl<'src, 'arena> {
    /// var(<...>) e.g. var(Display, "Advanced")
    /// Each item is an `ArenaNode`, so token locations are preserved.
    pub paren_specs: Option<ArenaVec<'arena, VarEditorSpecifierRef<'src, 'arena>>>,

    /// variable modifiers like public/protected/private/static/const/...
    /// Each modifier is an `ArenaNode` capturing its span; order preserved.
    pub modifiers: ArenaVec<'arena, VarModifier>,

    pub type_spec: TypeSpecifierRef<'src, 'arena>, // Named/InlineEnum/InlineStruct
    pub declarators: ArenaVec<'arena, VariableDeclaratorRef<'src, 'arena>>, // a, b=expr
    pub span: AstSpan,
}
pub type ClassVarDeclRef<'src, 'arena> = ArenaNode<'arena, ClassVarDecl<'src, 'arena>>;

#[derive(Debug)]
pub struct ClassConstDecl<'src, 'arena> {
    pub name: IdentifierToken,
    pub value: DeclarationLiteralRef<'src, 'arena>,
    pub span: AstSpan,
}
pub type ClassConstDeclRef<'src, 'arena> = ArenaNode<'arena, ClassConstDecl<'src, 'arena>>;

pub enum ClassMember<'src, 'arena>
where
    'src: 'arena,
{
    Function(CallableDefinitionRef<'src, 'arena>),
    TypeDefEnum(EnumDefRef<'src, 'arena>),
    TypeDefStruct(StructDefRef<'src, 'arena>),
    Var(ClassVarDeclRef<'src, 'arena>),

    Replication(ReplicationBlockRef<'src, 'arena>),
    State(StateDeclRef<'src, 'arena>),
    Const(ClassConstDeclRef<'src, 'arena>),
    Exec(ExecDirectiveRef<'arena>),
}

pub type ClassMemberRef<'src, 'arena> = ArenaNode<'arena, ClassMember<'src, 'arena>>;

#[derive(Clone, Copy, Debug)]
pub enum Reliability {
    Reliable,
    Unreliable,
}

#[derive(Debug)]
pub struct ReplicationRule<'src, 'arena> {
    pub reliability: Reliability, // reliable|unreliable
    pub condition: Option<ExpressionRef<'src, 'arena>>, // if (<expr>) or None
    pub members: ArenaVec<'arena, IdentifierToken>, // a, b, Foo()
    pub span: AstSpan,
}
pub type ReplicationRuleRef<'src, 'arena> = ArenaNode<'arena, ReplicationRule<'src, 'arena>>;

#[derive(Debug)]
pub struct ReplicationBlock<'src, 'arena> {
    pub rules: ArenaVec<'arena, ReplicationRuleRef<'src, 'arena>>,
    pub span: AstSpan,
}
pub type ReplicationBlockRef<'src, 'arena> = ArenaNode<'arena, ReplicationBlock<'src, 'arena>>;

// ---------- States ----------

#[derive(Clone, Copy, Debug)]
pub enum StateModifier {
    Auto,      // 'auto'
    Simulated, // 'simulated'
}

#[derive(Debug)]
pub struct StateDecl<'src, 'arena> {
    pub name: IdentifierToken,
    pub parent: Option<IdentifierToken>, // 'extends BaseState'
    pub modifiers: ArenaVec<'arena, StateModifier>, // auto, simulated
    pub ignores: Option<ArenaVec<'arena, IdentifierToken>>, // 'ignores Foo, Bar;'
    /// Body: ordinary statements plus nested function definitions (see `Statement::Function`).
    pub body: ArenaVec<'arena, StatementRef<'src, 'arena>>,
    pub span: AstSpan,
}
pub type StateDeclRef<'src, 'arena> = ArenaNode<'arena, StateDecl<'src, 'arena>>;

// NEW: exec directive node
#[derive(Debug)]
pub struct ExecDirective<'arena> {
    pub text: ArenaString<'arena>, // full line without trailing newline(s)
    pub span: AstSpan,
}
pub type ExecDirectiveRef<'arena> = ArenaNode<'arena, ExecDirective<'arena>>;

/// Keep your existing `ClassDeclaration` as the header.
/// Optionally: `pub type ClassHeader<'src, 'arena> = ClassDeclaration<'src, 'arena>;`
pub struct ClassDefinition<'src, 'arena>
where
    'src: 'arena,
{
    pub header: ClassDeclaration<'arena>, // or ClassHeader if you rename
    pub members: ArenaVec<'arena, ClassMemberRef<'src, 'arena>>,
}
