"""Time utilities: timestamp conversion, Shanghai timezone helpers."""

from __future__ import annotations

from datetime import datetime, timezone
from zoneinfo import ZoneInfo

_SH_TZ = ZoneInfo("Asia/Shanghai")


def ts_ms_now() -> int:
    """Return the current UTC timestamp in milliseconds."""
    return int(datetime.now(timezone.utc).timestamp() * 1000)


def ts_ms_to_shanghai_str(ts_ms: int, *, ms: bool = True) -> str:
    """Convert a UTC millisecond timestamp to a Shanghai-local datetime string.

    Parameters
    ----------
    ts_ms : int
        Millisecond epoch timestamp.
    ms : bool
        If True, include milliseconds (default: True).

    Returns
    -------
    str
        e.g. "2026-03-25 14:30:00.123"
    """
    dt = datetime.fromtimestamp(int(ts_ms) / 1000, tz=_SH_TZ)
    if ms:
        return dt.strftime("%Y-%m-%d %H:%M:%S.%f")[:-3]
    return dt.strftime("%Y-%m-%d %H:%M:%S")


def shanghai_str_to_ts_ms(s: str) -> int:
    """Parse a Shanghai-local datetime string to a UTC millisecond timestamp.

    Supports formats: "2026-03-25 14:30:00.123", "2026-03-25 14:30:00", "2026-03-25".

    Parameters
    ----------
    s : str
        Shanghai-local datetime string.

    Returns
    -------
    int
        UTC millisecond epoch timestamp.
    """
    s = s.strip()
    for fmt in ("%Y-%m-%d %H:%M:%S.%f", "%Y-%m-%d %H:%M:%S", "%Y-%m-%d"):
        try:
            dt = datetime.strptime(s, fmt).replace(tzinfo=_SH_TZ)
            return int(dt.timestamp() * 1000)
        except ValueError:
            continue
    raise ValueError(f"Cannot parse datetime string: {s!r}")


def now_shanghai_dt() -> datetime:
    """Return the current Shanghai datetime as a timezone-aware object."""
    return datetime.now(_SH_TZ)


def now_shanghai_str(*, ms: bool = True) -> str:
    """Return the current Shanghai time as a string."""
    return ts_ms_to_shanghai_str(ts_ms_now(), ms=ms)
