"""Context metadata."""

from enum import StrEnum, auto, nonmember
from typing import ClassVar, Self

from astropydantic import AstroPydanticTime, AstroPydanticUnit
from pydantic import model_validator

from ..._core import AstroField, ModelBase, doc, enum_doc
from .. import sciops

_NAMESPACE = "CTAO.DataProducts.Context"

__all__ = [
    "EnergyCoverage",
    "TimeRefPos",
    "TimeSystem",
    "TimeCoverage",
    "SpaceCoverage",
    "GeodeticEarthLocation",
    "SpatialFrame",
    "Observation",
    "Atmosphere",
    "TrackingCoverage",
]


class EnergyCoverage(ModelBase):
    """Describes spectral energy coverage of a data product."""

    _namespace: ClassVar[str] = _NAMESPACE

    energy_unit: AstroPydanticUnit = AstroField(
        "Unit for all energy metadata.", fits_keyword="EUNIT"
    )
    energy_min: float = AstroField(
        description=(
            "Approximate minimum energy of the dataset, precise enough for discovering"
            " the data product by energy range. Note that this maps to the IVOA ObsCore"
            " ``em_min`` keyword, however that is expressed as a wavelength."
        ),
        unit="TeV",
        ucd="em.energy;stat.min",
        ivoa_keyword="",
        fits_keyword="EMIN",
    )

    energy_max: float = AstroField(
        description=(
            "Approximate maximum energy of the dataset, precise enough for discovering"
            " the data product by energy range. Note that this maps to the IVOA ObsCore"
            " ``em_max`` keyword, however that is expressed as a wavelength."
        ),
        unit="TeV",
        ucd="em.energy;stat.max",
        fits_keyword="EMAX",
    )


class TimeRefPos(StrEnum):
    """
    Spatial location at which the observation time is valid.

    From FITS standard v4
    """

    _namespace = nonmember(_NAMESPACE)
    TOPOCENTER = "TOPOCENTER"
    RELOCATABLE = "RELOCATABLE"


class TimeSystem(StrEnum):
    """The time scale of the time-related keywords.

    For simulated data, use LOCAL.
    """

    _namespace = nonmember(_NAMESPACE)
    TT = auto()  #: Terrestrial Time
    UTC = auto()  #: Coordinated Universal Time
    UT1 = auto()  #: Universal Time
    TAI = auto()  #: International Atomic Time
    GPS = auto()  #: Global Positioning System time
    LOCAL = auto()  #: Local Time


class TimeReference(ModelBase):
    """
    Reference time for floating point time representations.

    For timestamps represented as ISO time strings, this is not relevant.
    """

    _namespace: ClassVar[str] = _NAMESPACE

    position: TimeRefPos = AstroField(enum_doc(TimeRefPos), fits_keyword="TREFPOS")
    time_mjd: float = AstroField(
        "Reference time in MJD", unit="day", fits_keyword="MJDREF"
    )
    unit: AstroPydanticUnit = AstroField(default="s", fits_keyword="TIMEUNIT")
    system: TimeSystem = AstroField(
        enum_doc(TimeSystem), default=TimeSystem.TAI, fits_keyword="TIMESYS"
    )
    resolution: float | None = AstroField(
        "time resolution in the time unit.",
        default=None,
        fits_keyword="TIMEDELT",
        ivoa_keyword="t_resolution",
    )


class TimeCoverage(ModelBase):
    """Time bounds for discoverability of a data product with time coverage."""

    _namespace: ClassVar[str] = _NAMESPACE

    reference: TimeReference | None = AstroField(
        description=doc(TimeReference), default=None
    )
    t_min: float | AstroPydanticTime = AstroField(
        "Start of time range of the data as either an ISO string or a float in the"
        " system defined by the time reference.",
        fits_keyword="TSTART",
    )
    t_max: float | AstroPydanticTime = AstroField(
        "End of time range of the data, as either an ISO string or a float in the"
        " system defined by the time reference.",
        fits_keyword="TSTOP",
    )

    @model_validator(mode="after")
    def _check_if_reference_is_needed(self) -> Self:
        """Validate TimeCoverage."""
        if isinstance(self.t_min, float) or isinstance(self.t_max, float):
            if self.reference is None:
                raise ValueError(
                    "if t_min/t_max are floats, you must specify a TimeReference."
                )
        return self


class SpatialFrame(StrEnum):
    """
    Standard equatorial coordinate system used for RA/Dec coordinates.

    Note that for tablular data products, columns of coordinate values have
    their own standards in FITS and IVOA for specifying the frame, e.g. the
    ``RADESYSn`` column metadata attribute.
    """

    _namespace = nonmember(_NAMESPACE)
    ICRS = "ICRS"  #: only ICRS is currently supported.


class SpaceCoverage(ModelBase):
    """
    Defines spatial coverage of a data product.

    The circle defined by ra, dec, field_of_view is used for the discoverability
    of data products by cone search, and should be considered a bounding circle
    if the true region covered is non-circular. The optional fields
    region_of_interest and moc can be used for more precise definition of the
    coverage.
    """

    _namespace: ClassVar[str] = _NAMESPACE

    frame: SpatialFrame = AstroField(
        description=doc(SpatialFrame), default=SpatialFrame.ICRS, fits_keyword="RADESYS"
    )

    ra: float = AstroField(
        description="ICRS Right ascension of the center of the region covered.",
        unit="deg",
        fits_keyword="RA_PNT",
        ivoa_keyword="s_ra",
        ucd="pos.eq.ra",
    )

    dec: float = AstroField(
        description="ICRS Declination of the center of the region covered",
        unit="deg",
        fits_keyword="DEC_PNT",
        ivoa_keyword="s_dec",
        ucd="pos.eq.dec",
    )

    field_of_view: float = AstroField(
        description=(
            "Approximate diameter (not radius) of the region covered by this data"
            " product if approximated as a circle on a sphere. For energy-dependent"
            " FOVs, a representative value should be chosen such that the data product"
            " can be identified in a cone search. Note that more accurate ROI shapes"
            " can be specified using the `region_of_interest` field, but"
            " `field_of_view` should always be filled because it provides a simple way"
            " to do a rough filtering."
        ),
        unit="deg",
        fits_keyword="FOV",
        ivoa_keyword="s_fov",
        ucd="pos.spherical.r;instr.fov",
    )

    region_of_interest: str | None = AstroField(
        default=None,
        description="IVOA ADQL-format region string describing the region of interest.",
        ivoa_keyword="s_region",
        fits_keyword="REGION",
        examples=[
            "Box ICRS 83.633083 22.0145 5 5",
            "Circle ICRS 83.633083 22.0145 5",
            (
                "Polygon ICRS 81.133083 19.5145  86.133083 19.5145  86.133083 24.5145 "
                " 81.133083 24.5145"
            ),
        ],
        ucd="meta.coverage;obs.field",
    )

    moc: str | None = AstroField(
        description=(
            "Multi-order coverage (MOC) string describing the region "
            "covered by the data product. https://www.ivoa.net/documents/MOC/"
        ),
        default=None,
        examples=[
            "4/2609 2611-2612  5/10350-10351 10362-10363 10433-10435 10440-10441 10443"
            " 10452 10456-10458   10468-10469  6/41398-41399 41442-41443 41446-41447"
            " 41464 41466-41467 41769-41771 41816-41818   41836 41840 41857 41860-41861"
            " 41863 41920"
        ],
        ucd="meta.coverage;obs.field",
    )


class TrackingCoverage(ModelBase):
    """
    Describes the coverage of ground-referenced observing parameters.

    This is required only for data products where knowing the horizontal (ground
    based alt/az) coverage is relevant. For equatorial tracking, elevation and
    azimuth vary with time. For fixed-horizontal tracking, ranges may collapse
    to single values.
    """

    _namespace: ClassVar[str] = _NAMESPACE

    elevation_min: float = AstroField(
        "Minimum elevation above horizon of observation.",
        unit="deg",
        ucd="pos.az.alt;stat.min",
        fits_keyword="ALT_MIN",
    )
    elevation_max: float = AstroField(
        "Maximum elevation above horizon of observation.",
        unit="deg",
        ucd="pos.az.alt;stat.max",
        fits_keyword="ALT_MAX",
    )
    azimuth_mean: float = AstroField(
        "Mean azimith of observation.",
        unit="deg",
        ucd="pos.az.azi;stat.mean",
        fits_keyword="AZ_MEAN",
    )
    azimuth_range: float = AstroField(
        "Full angular width of the observation in azimuth. Defines the horizontal"
        " coverage around `azimuth_mean`. The value should be capped at 360 deg, even"
        " if the telecope moves for larger than one rotation during the observation.",
        ge=0.0,
        le=360.0,
        unit="deg",
        ucd="pos.az.azi;stat.interval",
        fits_keyword="AZ_RANGE",
    )
    mode: sciops.TrackingMode = AstroField(
        doc(sciops.TrackingMode), fits_keyword="TRACKING"
    )


class Coverage(ModelBase):
    """
    Describes the spatial, temporal, spectral, and tracking coverage of a data product.

    This information is for the discovery of data products, not to specify
    precise bounds for science analysis. Typically this information can only be
    attached to higher-level data products, e.g. DL3 and above, as the
    information may not be available before processing.
    """

    _namespace: ClassVar[str] = _NAMESPACE

    time: TimeCoverage | None = AstroField(description=doc(TimeCoverage), default=None)
    space: SpaceCoverage | None = AstroField(
        description=doc(SpaceCoverage), default=None
    )
    energy: EnergyCoverage | None = AstroField(
        description=doc(EnergyCoverage), default=None
    )
    tracking: TrackingCoverage | None = AstroField(
        description=doc(TrackingCoverage), default=None
    )


class GeodeticEarthLocation(ModelBase):
    """Position on Earth of the observatory in geodetic coordinates."""

    # TODO: allow serialization to "classic" keywords "GEOLAT, GEOLON, ALTITUDE."

    _namespace: ClassVar[str] = _NAMESPACE

    longitude: float = AstroField(
        "Longitude, with East positive.", unit="deg", fits_keyword="OBSGEO-L"
    )
    latitude: float = AstroField(
        "Latitude, with North positive.", unit="deg", fits_keyword="OBSGEO-B"
    )
    height: float = AstroField(
        "Height above sea level", unit="m", fits_keyword="OBSGEO-H"
    )


class Atmosphere(ModelBase):
    """Overall atmospheric quality of the observation."""

    _namespace: ClassVar[str] = _NAMESPACE

    quality: sciops.SkyQuality = AstroField(
        enum_doc(sciops.SkyQuality), fits_keyword="SKYQUAL"
    )
    brightness: sciops.SkyBrightness = AstroField(
        enum_doc(sciops.SkyBrightness), fits_keyword="SKYBRIGH"
    )


class Observation(ModelBase):
    """Context info for data products associated with an observation."""

    _namespace: ClassVar[str] = _NAMESPACE

    coverage: Coverage | None = AstroField(doc(Coverage), default=None)
    location: GeodeticEarthLocation | None = AstroField(
        doc(GeodeticEarthLocation), default=None
    )
