"""Point-in-time dividend calendar helpers for leakage-safe ex-div scheduling."""
from __future__ import annotations

import datetime as dt
from bisect import bisect_right
from dataclasses import dataclass
from functools import lru_cache
from pathlib import Path
from typing import Any, Iterable

import pyarrow.parquet as pq

from libs.common.logging import get_logger

logger = get_logger(__name__)


@dataclass(frozen=True)
class DividendCalendarEntry:
    symbol: str
    as_of_date: dt.date
    ex_dividend_date: dt.date
    amount: float
    declaration_date: dt.date | None = None
    record_date: dt.date | None = None
    payment_date: dt.date | None = None
    source: str | None = None


def _coerce_date(value: Any) -> dt.date | None:
    if value is None:
        return None
    if isinstance(value, dt.datetime):
        return value.date()
    if isinstance(value, dt.date):
        return value
    text = str(value).strip()
    if not text:
        return None
    try:
        return dt.date.fromisoformat(text[:10])
    except ValueError:
        return None


def _coerce_float(value: Any) -> float | None:
    if value is None:
        return None
    try:
        return float(value)
    except (TypeError, ValueError):
        return None


class PointInTimeDividendCalendar:
    """Latest-known dividend schedule for each symbol as of a prior date."""

    def __init__(self, entries: Iterable[DividendCalendarEntry]) -> None:
        grouped: dict[str, list[DividendCalendarEntry]] = {}
        for entry in entries:
            grouped.setdefault(entry.symbol, []).append(entry)
        self._entries_by_symbol: dict[str, tuple[list[dt.date], list[DividendCalendarEntry]]] = {}
        for symbol, symbol_entries in grouped.items():
            symbol_entries.sort(key=lambda entry: (entry.as_of_date, entry.ex_dividend_date))
            as_of_dates = [entry.as_of_date for entry in symbol_entries]
            self._entries_by_symbol[symbol] = (as_of_dates, symbol_entries)

    @classmethod
    def from_parquet(cls, path: Path) -> PointInTimeDividendCalendar:
        table = pq.read_table(str(path))
        rows = table.to_pylist()
        entries: list[DividendCalendarEntry] = []
        for row in rows:
            symbol = str(row.get("symbol") or row.get("ticker") or "").strip().upper()
            if not symbol:
                continue
            as_of_date = _coerce_date(
                row.get("as_of_date")
                or row.get("declaration_date")
                or row.get("known_as_of_date")
            )
            ex_dividend_date = _coerce_date(
                row.get("ex_dividend_date")
                or row.get("ex_date")
            )
            amount = _coerce_float(row.get("amount") or row.get("dividend_amount"))
            if as_of_date is None or ex_dividend_date is None or amount is None:
                continue
            entries.append(
                DividendCalendarEntry(
                    symbol=symbol,
                    as_of_date=as_of_date,
                    ex_dividend_date=ex_dividend_date,
                    amount=amount,
                    declaration_date=_coerce_date(row.get("declaration_date")),
                    record_date=_coerce_date(row.get("record_date")),
                    payment_date=_coerce_date(row.get("payment_date")),
                    source=str(row.get("source")) if row.get("source") is not None else None,
                )
            )
        logger.info(
            "pit_dividend_calendar_loaded",
            path=str(path),
            rows=len(entries),
            symbols=len({entry.symbol for entry in entries}),
        )
        return cls(entries)

    def get_known_upcoming_ex_dividends(
        self,
        as_of_date: dt.date,
        allowed_ex_dates: Iterable[dt.date],
        symbols: Iterable[str] | None = None,
    ) -> dict[str, DividendCalendarEntry]:
        allowed = set(allowed_ex_dates)
        if not allowed:
            return {}
        symbol_filter = {
            str(symbol).strip().upper()
            for symbol in (symbols or [])
            if str(symbol).strip()
        }
        result: dict[str, DividendCalendarEntry] = {}
        for symbol, (as_of_dates, entries) in self._entries_by_symbol.items():
            if symbol_filter and symbol not in symbol_filter:
                continue
            idx = bisect_right(as_of_dates, as_of_date) - 1
            if idx < 0:
                continue
            entry = entries[idx]
            if entry.ex_dividend_date in allowed:
                result[symbol] = entry
        return result


@lru_cache(maxsize=8)
def load_pit_dividend_calendar(path_str: str) -> PointInTimeDividendCalendar | None:
    path = Path(path_str)
    if not path.exists():
        logger.info("pit_dividend_calendar_missing", path=str(path))
        return None
    return PointInTimeDividendCalendar.from_parquet(path)


__all__ = [
    "DividendCalendarEntry",
    "PointInTimeDividendCalendar",
    "load_pit_dividend_calendar",
]