from enum import StrEnum
from typing import Annotated, Literal, Self

from pydantic import BaseModel, Field, model_validator
from typing_extensions import TypeAliasType

MAX_TIMESTEPS = 12  # OlmoEarth's max number of timesteps


class Month(StrEnum):
    JANUARY = "January"
    FEBRUARY = "February"
    MARCH = "March"
    APRIL = "April"
    MAY = "May"
    JUNE = "June"
    JULY = "July"
    AUGUST = "August"
    SEPTEMBER = "September"
    OCTOBER = "October"
    NOVEMBER = "November"
    DECEMBER = "December"


class ObservationStrategyName(StrEnum):
    SAMPLED = "sampled"
    REPEATING_INTERVAL = "repeating_interval"


class TemporalityName(StrEnum):
    PERIOD = "period"
    POINT_IN_TIME = "point_in_time"
    FRAMED_POINT_IN_TIME = "framed_point_in_time"


class SampledObservation(BaseModel):
    """Collect discrete snapshots within a window."""

    name: Literal[ObservationStrategyName.SAMPLED]
    duration_days: int = Field(ge=1, description="The observation window duration.")
    num_samples: int = Field(
        ge=1, le=MAX_TIMESTEPS, description="Number of samples to collect."
    )


class RepeatingIntervalObservation(BaseModel):
    """Divide window into fixed sub-periods, one observation (mosaic or snapshot) per period."""

    name: Literal[ObservationStrategyName.REPEATING_INTERVAL]
    num_periods: int = Field(
        ge=1, le=MAX_TIMESTEPS, description="Number of intervals."
    )
    period_duration_days: int = Field(ge=1, description="Duration of each interval.")


ObservationStrategy = TypeAliasType(
    "ObservationStrategy",
    Annotated[
        SampledObservation | RepeatingIntervalObservation, Field(discriminator="name")
    ],
)


class PeriodTemporality(BaseModel):
    """
    Represents the duration of time we're making a prediction for.
    The observation window equals the prediction period.
    request.time_of_interest marks the start of the prediction and observation periods.
    """

    name: Literal[TemporalityName.PERIOD]
    observation: ObservationStrategy
    allowed_start_month_constraints: set[Month] | None = Field(
        default=None, description="Months valid for starting this period."
    )


class PointInTimeTemporality(BaseModel):
    """
    Represents a specific instant we're making a prediction for.
    Uses a narrow symmetric window centered on request.time_of_interest to find imagery.
    """

    name: Literal[TemporalityName.POINT_IN_TIME]
    observation_window_seconds: int = Field(
        ge=1,
        description="Total time covered by the symmetric observation window centered on the prediction point.",
    )


class FramedPointInTimeTemporality(BaseModel):
    """
    Represents a specific instant we're making a prediction for, with context from before and/or after.
    Optional offsets allow gaps between the prediction point and observation windows.
    """

    name: Literal[TemporalityName.FRAMED_POINT_IN_TIME]
    lookbehind_observation: ObservationStrategy | None = Field(
        default=None, description="Observation window preceding the prediction point."
    )
    lookbehind_offset_days: int | None = Field(
        default=None,
        ge=0,
        description="Gap between end of lookbehind window and prediction point.",
    )
    lookahead_observation: ObservationStrategy | None = Field(
        default=None, description="Observation window following the prediction point."
    )
    lookahead_offset_days: int | None = Field(
        default=None,
        ge=0,
        description="Gap between prediction point and start of lookahead window.",
    )

    @model_validator(mode="after")
    def at_least_one_observation(self) -> Self:
        if self.lookbehind_observation is None and self.lookahead_observation is None:
            raise ValueError(
                "At least one of lookbehind_observation or lookahead_observation must be set"
            )
        return self

    @model_validator(mode="after")
    def offsets_require_observations(self) -> Self:
        if (
            self.lookbehind_offset_days is not None
            and self.lookbehind_observation is None
        ):
            raise ValueError(
                "lookbehind_offset_days requires lookbehind_observation to be set"
            )
        if (
            self.lookahead_offset_days is not None
            and self.lookahead_observation is None
        ):
            raise ValueError(
                "lookahead_offset_days requires lookahead_observation to be set"
            )
        return self


Temporality = TypeAliasType(
    "Temporality",
    Annotated[
        PeriodTemporality | PointInTimeTemporality | FramedPointInTimeTemporality,
        Field(discriminator="name"),
    ],
)
