"""Core classes for the data model."""

from enum import Flag, StrEnum
from inspect import cleandoc
from typing import Any

from astropy import units as u
from astropy.io.votable.ucd import check_ucd, parse_ucd
from pydantic import (
    AfterValidator,
    BaseModel,
    ConfigDict,
    Field,
    validate_call,
)
from pydantic_core import core_schema

__all__ = [
    "AstroField",
    "ModelBase",
    "doc",
    "enum_doc",
    "ValidUCD",
    "type_to_string",
    "StrFlag",
]


def is_valid_ucd(value: str) -> str:
    if check_ucd(value, check_controlled_vocabulary=True):
        return value
    msg = f"Invalid UCD: '{value}'"
    raise ValueError(msg)


ValidUCD = AfterValidator(is_valid_ucd)


@validate_call
def AstroField(  # noqa: N802
    description: str | None = None,
    *,
    fits_keyword: str | None = None,
    ivoa_keyword: str | None = None,
    unit: str | None = None,
    ucd: str | None = None,
    examples: list[str] | None = None,
    **kwargs,
) -> Field:  # type: ignore
    """
    Return a Field with extra astronomy-related metadata.

    The extra info is not for validation, but used for documentation or
    serialization. This is just a helper to avoid having to add a
    json_schema_extra dict manually.
    """
    if fits_keyword and len(fits_keyword) > 8:
        msg = f"FITS Keyword muse be less than 8 characters, got '{fits_keyword}'"
        raise ValueError(
            msg,
        )

    # normalize unit string
    if unit:
        unit = u.Unit(unit).to_string("fits")

    json_schema_extra: dict[str, Any] = dict()
    if fits_keyword:
        json_schema_extra["fits_keyword"] = fits_keyword
    if ivoa_keyword:
        json_schema_extra["ivoa_keyword"] = ivoa_keyword
    if unit:
        json_schema_extra["unit"] = unit
    if examples:
        json_schema_extra["examples"] = examples
    if ucd:
        parse_ucd(ucd, check_controlled_vocabulary=True)
        json_schema_extra["ucd"] = ucd

    return Field(description=description, json_schema_extra=json_schema_extra, **kwargs)


def doc(obj):
    """Return cleaned up docstring for the object, suitable for a description field."""
    doc = obj.__doc__ or ""
    return cleandoc(doc).replace("\n", " ")


def enum_doc(enum: type[StrEnum]) -> str:
    """Return nicer documentation for an Enum."""
    return doc(enum) + " Options are: " + ", ".join([f'"{v}"' for v in enum]) + "."


class ModelBase(BaseModel):
    """
    Base class for all CTAO models.

    Should just set the common model_config here.
    """

    model_config = ConfigDict(extra="forbid")


def type_to_string(cls):
    """Turn a type into a string."""
    if hasattr(cls, "__args__"):
        # Handling compound types
        return f"{'|'.join(map(type_to_string, cls.__args__))}"
    elif hasattr(cls, "__name__"):
        return cls.__name__
    elif hasattr(cls, "__forward_arg__"):
        # Handle forward refs
        return cls.__forward_arg__
    return str(cls)


class StrFlag(Flag):
    """Like enum.Flag, but allowing string input, like ``A|B``."""

    @classmethod
    def _missing_(cls, value):
        # Only handle strings; defer everything else to Flag
        if not isinstance(value, str):
            return super()._missing_(value)

        result = cls(0)
        parts = value.split("|")

        for part in parts:
            name = part.strip().upper()
            try:
                result |= cls[name]
            except KeyError:
                raise ValueError(
                    f"{name!r} is not a valid member of {cls.__name__}"
                ) from None

        return result

    def to_string(self) -> str:
        """Turn the flag into a string representation."""
        if self.value == 0:
            return "0"
        return "|".join(member.name for member in type(self) if member in self)

    def __str__(self):
        return self.to_string()

    @classmethod
    def __get_pydantic_core_schema__(cls, source, handler):
        """Implement correct serialization for pydantic."""
        return core_schema.no_info_after_validator_function(
            cls,
            core_schema.union_schema(
                [
                    core_schema.is_instance_schema(cls),
                    core_schema.str_schema(),
                    core_schema.int_schema(),
                ]
            ),
            serialization=core_schema.plain_serializer_function_ser_schema(
                lambda v: v.to_string(),
                return_schema=core_schema.str_schema(),
            ),
        )
