"""Data models for the Open Energy ID."""

import datetime as dt
from typing import Self, overload

import pandas as pd
import polars as pl
from pydantic import BaseModel


class TimeSeriesBase(BaseModel):
    """Pydantic base model for time series data."""

    index: list[dt.datetime]

    @classmethod
    def from_pandas(cls, data: pd.Series | pd.DataFrame) -> Self:
        """Create from a Pandas Object."""
        raise NotImplementedError

    def to_pandas(self, timezone: str = "UTC") -> pd.Series | pd.DataFrame:
        """Convert to a Pandas Object."""
        raise NotImplementedError

    @overload
    def to_json(self, path: None = None, **kwargs) -> str:
        """Dump to a JSON string."""

    @overload
    def to_json(self, path: str, **kwargs) -> None:
        """Dump to a JSON file."""

    def to_json(self, path: str | None = None, **kwargs) -> str | None:
        """Dump to a JSON string or file."""
        if path is None:
            return self.model_dump_json(**kwargs)
        encoding = kwargs.pop("encoding", "UTF-8")
        with open(path, "w", encoding=encoding) as file:
            file.write(self.model_dump_json(**kwargs))
        return None

    @overload
    @classmethod
    def from_json(cls, string: str, **kwargs) -> Self:
        """Load from a JSON string."""

    @overload
    @classmethod
    def from_json(cls, *, path: str, **kwargs) -> Self:
        """Load from a JSON file."""

    @classmethod
    def from_json(cls, string: str | None = None, path: str | None = None, **kwargs) -> Self:
        """Load from a JSON file or string."""
        if string:
            return cls.model_validate_json(string, **kwargs)
        if path:
            encoding = kwargs.pop("encoding", "UTF-8")
            with open(path, encoding=encoding) as file:
                return cls.model_validate_json(file.read(), **kwargs)
        raise ValueError("Either string or path must be provided.")

    def first_timestamp(self) -> dt.datetime:
        """Get the first timestamp in the index."""
        return min(self.index)

    def last_timestamp(self) -> dt.datetime:
        """Get the last timestamp in the index."""
        return max(self.index)


class TimeSeries(TimeSeriesBase):
    """
    Represents a time series data.
    Attributes:
        name (str | None): The name of the time series.
        data (list[float | None]): The data points of the time series.
    Methods:
        replace_nan_with_none(cls, data: list[float]) -> list[float | None]:
            Replace NaN values with None.
        from_pandas(cls, data: pd.Series) -> Self:
            Create a TimeSeries object from a Pandas Series.
        to_pandas(self, timezone: str = "UTC") -> pd.Series:
            Convert the TimeSeries object to a Pandas Series.
        from_polars(cls, data: pl.DataFrame | pl.LazyFrame) -> Self:
            Create a TimeSeries object from Polars data.
        to_polars(self, timezone: str = "UTC") -> pl.LazyFrame:
            Convert the TimeSeries object to a Polars LazyFrame.
    """

    name: str | None = None
    data: list[float | None]

    @classmethod
    def from_pandas(cls, data: pd.Series) -> Self:
        """Create from a Pandas Series."""
        if not hasattr(data, "name") or data.name is None:
            name = None
        else:
            name = str(data.name)
        return cls(name=name, data=data.tolist(), index=data.index.tolist())

    def to_pandas(self, timezone: str = "UTC") -> pd.Series:
        """Convert to a Pandas Series with float dtype and None as NaN."""
        series = pd.Series(self.data, name=self.name, index=self.index, dtype=float)
        series.index = pd.to_datetime(series.index, utc=True)
        return series.tz_convert(timezone)

    @classmethod
    def from_polars(cls, data: pl.DataFrame | pl.LazyFrame) -> Self:
        """Create from Polars data."""
        # Always work with DataFrame
        df = data.collect() if isinstance(data, pl.LazyFrame) else data

        if len(df.columns) != 2:
            raise ValueError("Must contain exactly two columns: timestamp and value")

        value_col = [col for col in df.columns if col != "timestamp"][0]
        return cls(
            name=value_col,
            data=df[value_col].cast(pl.Float64).to_list(),  # Ensure float type
            index=df["timestamp"].cast(pl.Datetime).dt.convert_time_zone("UTC").to_list(),
        )

    def to_polars(self, timezone: str = "UTC") -> pl.LazyFrame:
        """Convert to Polars LazyFrame."""
        # Always return LazyFrame as specified in return type
        df = pl.DataFrame(
            {
                "timestamp": pl.Series(self.index, dtype=pl.Datetime).dt.convert_time_zone(
                    timezone
                ),
                "total" if self.name is None else self.name: pl.Series(self.data, dtype=pl.Float64),
            }
        )
        return df.lazy()


class TimeDataFrame(TimeSeriesBase):
    """Time series data with multiple columns."""

    columns: list[str]
    data: list[list[float | None]]

    @classmethod
    def from_pandas(cls, data: pd.DataFrame) -> Self:
        """Create from a Pandas DataFrame."""
        # Cast values to float | None
        values = [
            [float(x) if pd.notnull(x) else None for x in row] for row in data.values.tolist()
        ]
        return cls(columns=data.columns.tolist(), data=values, index=data.index.tolist())

    def to_pandas(self, timezone: str = "UTC") -> pd.DataFrame:
        """Convert to a Pandas DataFrame."""
        frame = pd.DataFrame(self.data, columns=self.columns, index=self.index)
        frame.index = pd.to_datetime(frame.index, utc=True)
        return frame.tz_convert(timezone)

    @classmethod
    def from_timeseries(cls, data: list[TimeSeries]) -> Self:
        """Create from a list of TimeSeries objects."""
        return cls(
            columns=[series.name or "" for series in data],  # Handle None names
            data=[series.data for series in data],  # Pass list of value lists
            index=data[0].index,
        )

    def to_timeseries(self) -> list[TimeSeries]:
        """Convert to a list of TimeSeries objects."""
        return [
            TimeSeries(name=col, data=[row[i] for row in self.data], index=self.index)
            for i, col in enumerate(self.columns)
        ]
