use crate::Db;
use ruff_db::files::{File, FileRange};
use ruff_db::parsed::parsed_module;
use ruff_db::source::source_text;
use ruff_text_size::{TextLen, TextRange};
use ty_module_resolver::Module;
use ty_python_core::definition::Definition;

#[derive(Debug, PartialEq, Eq, Hash)]
pub enum TypeDefinition<'db> {
    Module(Module<'db>),
    /// A class created via a `class` statement.
    StaticClass(Definition<'db>),
    /// A class created dynamically via `type(name, bases, dict)`.
    DynamicClass(Definition<'db>),
    Function(Definition<'db>),
    TypeVar(Definition<'db>),
    TypeAlias(Definition<'db>),
    NewType(Definition<'db>),
    SpecialForm(Definition<'db>),
    EnumMember(Definition<'db>),
}

impl TypeDefinition<'_> {
    pub fn focus_range(&self, db: &dyn Db) -> Option<FileRange> {
        match self {
            Self::Module(_) => None,
            Self::StaticClass(definition)
            | Self::DynamicClass(definition)
            | Self::Function(definition)
            | Self::TypeVar(definition)
            | Self::TypeAlias(definition)
            | Self::SpecialForm(definition)
            | Self::NewType(definition)
            | Self::EnumMember(definition) => {
                let module = parsed_module(db, definition.python_file(db)).load(db);
                Some(definition.focus_range(db, &module))
            }
        }
    }

    pub fn full_range(&self, db: &dyn Db) -> Option<FileRange> {
        match self {
            Self::Module(module) => {
                let file = module.file(db)?;
                let source = source_text(db, file);
                Some(FileRange::new(file, TextRange::up_to(source.text_len())))
            }
            Self::StaticClass(definition)
            | Self::DynamicClass(definition)
            | Self::Function(definition)
            | Self::TypeVar(definition)
            | Self::TypeAlias(definition)
            | Self::SpecialForm(definition)
            | Self::NewType(definition)
            | Self::EnumMember(definition) => {
                let module = parsed_module(db, definition.python_file(db)).load(db);
                Some(definition.full_range(db, &module))
            }
        }
    }

    pub(super) fn file(&self, db: &dyn Db) -> Option<File> {
        match self {
            Self::Module(module) => module.file(db),
            Self::StaticClass(definition)
            | Self::DynamicClass(definition)
            | Self::Function(definition)
            | Self::TypeVar(definition)
            | Self::TypeAlias(definition)
            | Self::SpecialForm(definition)
            | Self::NewType(definition)
            | Self::EnumMember(definition) => Some(definition.file(db)),
        }
    }
}

impl<'db> TypeDefinition<'db> {
    pub fn definition<'a>(&'a self) -> Option<&'a Definition<'db>> {
        match self {
            Self::Module(_) => None,
            Self::StaticClass(definition)
            | Self::DynamicClass(definition)
            | Self::Function(definition)
            | Self::TypeVar(definition)
            | Self::TypeAlias(definition)
            | Self::SpecialForm(definition)
            | Self::NewType(definition)
            | Self::EnumMember(definition) => Some(definition),
        }
    }
}
