"""Helpers for SDK-side energy-model query payload construction."""

from __future__ import annotations

from collections.abc import Mapping, Sequence
from copy import deepcopy
from datetime import UTC, date, datetime, timedelta, timezone
from zoneinfo import ZoneInfo

TimeStepFormat = str

_REQUIRED_QUERY_COLUMNS = ("date", "ghi", "dhi", "temperature", "windSpeed")
_OPTIONAL_QUERY_COLUMNS = ("dni", "relativeHumidity", "poa")
_COLUMN_ALIASES = {
    "timestamp": "date",
    "datetime": "date",
    "wind_speed": "windSpeed",
    "windspeed": "windSpeed",
    "temp": "temperature",
    "relative_humidity": "relativeHumidity",
    "relativehumidity": "relativeHumidity",
}


def build_query_payload(
    base_input: dict,
    weather_data: object,
    *,
    location: dict | None = None,
    start_date: object | None = None,
    end_date: object | None = None,
    time_step_format: TimeStepFormat = "utc",
    timezone_offset: float | None = None,
    timezone_name: str | None = None,
) -> dict:
    """Return an API-ready query payload with top-level inline ``weatherData``."""
    payload = deepcopy(base_input)
    if location is not None:
        payload.pop("locationId", None)
        payload.pop("projectId", None)
        payload["location"] = deepcopy(location)

    if "location" not in payload and "locationId" not in payload and "projectId" not in payload:
        raise ValueError("Query payload requires one of 'location', 'locationId', or 'projectId'.")

    resolved_offset, resolved_name = _resolve_time_context(
        payload,
        explicit_offset=timezone_offset,
        explicit_name=timezone_name,
        time_step_format=time_step_format,
    )
    payload["weatherData"] = prepare_query_weather_data(
        weather_data,
        start_date=start_date,
        end_date=end_date,
        time_step_format=time_step_format,
        timezone_offset=resolved_offset,
        timezone_name=resolved_name,
    )
    return payload


def prepare_query_weather_data(
    weather_data: object,
    *,
    start_date: object | None = None,
    end_date: object | None = None,
    time_step_format: TimeStepFormat = "utc",
    timezone_offset: float | None = None,
    timezone_name: str | None = None,
) -> dict[str, list]:
    """Normalize weather input into API-ready array payloads."""
    records = _normalize_records(weather_data)
    if not records:
        raise ValueError("Weather data is empty.")

    normalized_records = [_canonicalize_record(record) for record in records]
    _validate_required_columns(normalized_records)

    normalized_start = (
        _coerce_timestamp(start_date, time_step_format, timezone_offset=timezone_offset, timezone_name=timezone_name)[0]
        if start_date is not None
        else None
    )
    normalized_end = (
        _coerce_timestamp(end_date, time_step_format, timezone_offset=timezone_offset, timezone_name=timezone_name)[0]
        if end_date is not None
        else None
    )

    filtered_rows: list[tuple[dict, datetime]] = []
    for record in normalized_records:
        filter_dt, utc_dt = _coerce_timestamp(
            record["date"],
            time_step_format,
            timezone_offset=timezone_offset,
            timezone_name=timezone_name,
        )
        if normalized_start is not None and filter_dt < normalized_start:
            continue
        if normalized_end is not None and filter_dt > normalized_end:
            continue
        filtered_rows.append((record, utc_dt))

    if not filtered_rows:
        raise ValueError("Weather data is empty after applying the requested query window.")

    query_weather: dict[str, list] = {
        "date": [_to_epoch_millis(utc_dt) for _, utc_dt in filtered_rows],
    }
    for column in _REQUIRED_QUERY_COLUMNS:
        if column == "date":
            continue
        query_weather[column] = [record[column] for record, _ in filtered_rows]

    for column in _OPTIONAL_QUERY_COLUMNS:
        values = [record.get(column) for record, _ in filtered_rows]
        if all(value is not None for value in values):
            query_weather[column] = values

    return query_weather


def _resolve_time_context(
    payload: dict,
    *,
    explicit_offset: float | None,
    explicit_name: str | None,
    time_step_format: TimeStepFormat,
) -> tuple[float | None, str | None]:
    location = payload.get("location")
    location_offset = None
    location_name = None
    if isinstance(location, Mapping):
        location_offset = location.get("timeZone")
        location_name = location.get("timeZoneName")

    resolved_offset = explicit_offset if explicit_offset is not None else _coerce_optional_float(location_offset)
    resolved_name = explicit_name if explicit_name is not None else _coerce_optional_str(location_name)

    if time_step_format == "offset" and resolved_offset is None:
        raise ValueError("time_step_format='offset' requires timezone_offset or location.timeZone.")
    if time_step_format == "local" and resolved_name is None:
        raise ValueError("time_step_format='local' requires timezone_name or location.timeZoneName.")

    return resolved_offset, resolved_name


def _normalize_records(weather_data: object) -> list[dict]:
    if _is_pandas_dataframe(weather_data):
        return weather_data.to_dict(orient="records")

    if isinstance(weather_data, Mapping):
        return _records_from_columnar(weather_data)

    if (
        isinstance(weather_data, Sequence)
        and not isinstance(weather_data, (str, bytes, bytearray))
        and all(isinstance(item, Mapping) for item in weather_data)
    ):
        return [dict(item) for item in weather_data]

    raise ValueError("Weather data must be a pandas DataFrame, a list of records, or a dict of column arrays.")


def _records_from_columnar(columnar: Mapping) -> list[dict]:
    if "weatherData" in columnar and isinstance(columnar["weatherData"], Mapping):
        columnar = columnar["weatherData"]

    lengths = {
        key: len(value)
        for key, value in columnar.items()
        if isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray))
    }
    if not lengths:
        raise ValueError("Columnar weather data must contain sequence values.")
    if len(set(lengths.values())) != 1:
        raise ValueError(f"Weather data columns must have matching lengths. Got: {lengths}")

    row_count = next(iter(lengths.values()))
    records: list[dict] = []
    for idx in range(row_count):
        record = {}
        for key, value in columnar.items():
            if isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)):
                record[key] = value[idx]
        records.append(record)
    return records


def _canonicalize_record(record: Mapping) -> dict:
    normalized: dict[str, object] = {}
    for key, value in record.items():
        canonical_key = _COLUMN_ALIASES.get(str(key), str(key))
        normalized[canonical_key] = value
    return normalized


def _validate_required_columns(records: list[dict]) -> None:
    missing: list[str] = []
    for column in _REQUIRED_QUERY_COLUMNS:
        if any(record.get(column) is None for record in records):
            missing.append(column)
    if missing:
        raise ValueError(f"Weather data is missing required query columns: {', '.join(missing)}")


def _coerce_timestamp(
    value: object,
    time_step_format: TimeStepFormat,
    *,
    timezone_offset: float | None,
    timezone_name: str | None,
) -> tuple[datetime, datetime]:
    if time_step_format not in {"utc", "offset", "local"}:
        raise ValueError("time_step_format must be one of: utc, offset, local")

    dt_value = _parse_timestamp(value)
    if time_step_format == "utc":
        utc_dt = _to_utc_datetime(dt_value, UTC)
        return utc_dt, utc_dt

    if time_step_format == "offset":
        if timezone_offset is None:
            raise ValueError("Offset time-step formatting requires timezone_offset.")
        offset_tz = timezone(timedelta(hours=timezone_offset))
        offset_dt = _to_mode_datetime(dt_value, offset_tz)
        return offset_dt, offset_dt.astimezone(UTC)

    if timezone_name is None:
        raise ValueError("Local time-step formatting requires timezone_name.")
    local_tz = ZoneInfo(timezone_name)
    local_dt = _to_mode_datetime(dt_value, local_tz)
    return local_dt, local_dt.astimezone(UTC)


def _parse_timestamp(value: object) -> datetime:
    if isinstance(value, datetime):
        return value
    if isinstance(value, date):
        return datetime.combine(value, datetime.min.time())
    if isinstance(value, (int, float)):
        unit_divisor = 1000.0 if abs(value) >= 100_000_000_000 else 1.0
        return datetime.fromtimestamp(float(value) / unit_divisor, tz=UTC)
    if isinstance(value, str):
        normalized = value.strip()
        if normalized.endswith("Z"):
            normalized = normalized[:-1] + "+00:00"
        return datetime.fromisoformat(normalized)
    raise ValueError(f"Unsupported timestamp value: {value!r}")


def _to_utc_datetime(value: datetime, default_tz: timezone) -> datetime:
    if value.tzinfo is None:
        return value.replace(tzinfo=default_tz).astimezone(UTC)
    return value.astimezone(UTC)


def _to_mode_datetime(value: datetime, mode_tz: timezone | ZoneInfo) -> datetime:
    if value.tzinfo is None:
        return value.replace(tzinfo=mode_tz)
    return value.astimezone(mode_tz)


def _to_epoch_millis(value: datetime) -> int:
    return int(round(value.timestamp() * 1000))


def _is_pandas_dataframe(value: object) -> bool:
    cls = value.__class__
    return cls.__name__ == "DataFrame" and cls.__module__.startswith("pandas")


def _coerce_optional_float(value: object) -> float | None:
    if value is None:
        return None
    return float(value)


def _coerce_optional_str(value: object) -> str | None:
    if value is None:
        return None
    text = str(value).strip()
    return text or None
