"""KST buckets and session tags (pure Python).

Session tags (sentiment-labelling §6.2 hour bands; rollup column ``session``, classifier-plan §7.3):
``holiday`` on weekends and KRX holidays, else ``pre_open`` 06:00-09:00, ``regular`` 09:00-15:30,
``after_hours`` 15:30-20:00, ``night`` 20:00-06:00 (KST). An hourly bucket takes the tag of its start time, so
the 15:00 bucket (which holds the 15:30 close) is ``regular``. Holidays come from ``config/krx_holidays.yaml``
(a list of dates, or ``{"holidays": [...]}``); without it only weekends count.
"""
from __future__ import annotations

from datetime import date, datetime, time, timedelta, timezone
from pathlib import Path
from typing import Iterable

KST = timezone(timedelta(hours=9))
SESSIONS = ("pre_open", "regular", "after_hours", "night", "holiday")


def to_kst(dt: datetime) -> datetime:
    if dt.tzinfo is None:
        raise ValueError("naive datetime: written_at must carry a zone")
    return dt.astimezone(KST)


def parse(text: str) -> datetime:
    """ISO time with an offset (``Z`` accepted) -> aware datetime."""
    dt = datetime.fromisoformat(text.replace("Z", "+00:00"))
    if dt.tzinfo is None:
        raise ValueError(f"time without offset: {text!r}")
    return dt


def hour_bucket(dt: datetime) -> datetime:
    """KST hour start as a naive local timestamp (``bucket_kst timestamp``)."""
    k = to_kst(dt)
    return k.replace(minute=0, second=0, microsecond=0, tzinfo=None)


def day_bucket(dt: datetime) -> date:
    return to_kst(dt).date()


def is_trading_day(d: date, holidays: Iterable[date] = ()) -> bool:
    return d.weekday() < 5 and d not in set(holidays)


def session_tag(dt: datetime, holidays: Iterable[date] = ()) -> str:
    """Session of an aware time, or of a naive KST timestamp (a bucket start)."""
    k = dt.replace(tzinfo=None) if dt.tzinfo is None else to_kst(dt).replace(tzinfo=None)
    if not is_trading_day(k.date(), holidays):
        return "holiday"
    t = k.time()
    if time(6) <= t < time(9):
        return "pre_open"
    if time(9) <= t < time(15, 30):
        return "regular"
    if time(15, 30) <= t < time(20):
        return "after_hours"
    return "night"


def load_holidays(path: str | Path | None) -> set[date]:
    if path is None or not Path(path).exists():
        return set()
    import yaml

    value = yaml.safe_load(Path(path).read_text(encoding="utf-8")) or []
    if isinstance(value, dict):
        value = value.get("holidays", [])
    out = set()
    for v in value:
        if isinstance(v, dict):
            v = v.get("date")
        out.add(v if isinstance(v, date) else date.fromisoformat(str(v)))
    return out


def iso_week(dt: datetime) -> str:
    """ISO week of the KST date, e.g. ``2026-W42``."""
    y, w, _ = to_kst(dt).isocalendar()
    return f"{y}-W{w:02d}"


def load_groups(path: str | Path | None, codes: Iterable[str] = ()) -> dict[str, str]:
    """Code -> stock group/sector from ``config/stocks.yaml`` (classifier-plan §5.2).

    Accepted shapes: ``{"stocks": [{"code", "group"|"sector"}]}``, a list of such entries, or ``{code: group}``.
    Codes without an entry (or without the file) are their own group."""
    out = {c: c for c in codes}
    if path is None or not Path(path).exists():
        return out
    import yaml

    value = yaml.safe_load(Path(path).read_text(encoding="utf-8")) or {}
    entries = value.get("stocks", value) if isinstance(value, dict) else value
    if isinstance(entries, dict):
        out.update({str(k): str(v) for k, v in entries.items()})
        return out
    for e in entries:
        code = str(e["code"])
        out[code] = str(e.get("group") or e.get("sector") or code)
    return out