"""Definition of get_state function."""

from datetime import datetime, UTC
from pandera.typing import DataFrame
import pandas as pd
import pandera.pandas as pa
import sqlalchemy as sa

from utils.database.tables import (
    States,
    StatesMeta,
)
from utils.database.dataframes import SensorStateSchema
from utils.database import get_db_session


@pa.check_types
def get_state(
    entity_id: str,
    limit: int | None = None,
    start_time: datetime | None = None,
    end_time: datetime | None = None,
) -> DataFrame[SensorStateSchema]:
    """
    Retrieve sensor state data.

    Args:
        entity_id: The entity ID of the sensor.
        limit: Number of recent records to retrieve (default is None for all records).
        start_time: Start of time range (inclusive). If None, no lower bound.
        end_time: End of time range (inclusive). If None, no upper bound.

    Returns:
        DataFrame[SensorStateSchema]: The retrieved sensor state.
    """

    # Create a new database session
    session = get_db_session()

    # Prepare the base statement
    stmt = (
        sa.select(
            States.state,
            sa.func.to_timestamp(States.last_updated_ts).label("time"),
        )
        .join(
            StatesMeta,
            States.metadata_id == StatesMeta.metadata_id,
        )
        .where(StatesMeta.entity_id == entity_id)
    )

    # Add time range filters if provided
    if start_time is not None:
        start_ts = start_time.astimezone(UTC).timestamp()
        stmt = stmt.where(States.last_updated_ts >= start_ts)

    if end_time is not None:
        end_ts = end_time.astimezone(UTC).timestamp()
        stmt = stmt.where(States.last_updated_ts <= end_ts)

    # Order by time descending and apply limit if specified
    stmt = stmt.order_by(States.last_updated_ts.desc())

    if limit is not None:
        stmt = stmt.limit(limit)

    # Execute query
    with session.begin() as s:
        df = pd.read_sql(stmt, s.connection())

    # Handle empty results
    if len(df) == 0:
        # Create empty DataFrame with correct schema
        empty_df = pd.DataFrame(
            {
                "time": pd.Series([], dtype="datetime64[ns, UTC]"),
                "state": pd.Series([], dtype="object"),
            }
        )
        return DataFrame[SensorStateSchema](empty_df)

    # Reorder columns to match schema definition
    df = df[["time", "state"]]

    # Convert time column to nanosecond precision (PostgreSQL returns microseconds)
    df["time"] = df["time"].dt.as_unit("ns")

    return DataFrame[SensorStateSchema](df)