analytics.py

from __future__ import annotations

import hashlib
import json
import os
import threading
from dataclasses import dataclass
from datetime import UTC, datetime
from pathlib import Path
from typing import Literal

import pandas as pd
from huggingface_hub import download_bucket_files

BUCKET_ID = os.environ.get(
    "ANALYTICS_BUCKET",
    "wzsg/polymarket-orderfilled-analytics",
)
CACHE_DIR = Path(
    os.environ.get(
        "ANALYTICS_CACHE_DIR",
        "/tmp/polymarket-orderfilled-analytics",
    )
)
MONTHLY_FILE = "monthly_metrics.parquet"
DAILY_FILE = "daily_metrics.parquet"
STATE_FILE = "aggregation_state.json"

NUMERIC_COLUMNS = (
    "fill_event_count",
    "unique_transaction_count",
    "nominal_collateral_volume",
    "fee_amount",
)
Period = Literal["year", "quarter"]


@dataclass(frozen=True)
class Snapshot:
    monthly: pd.DataFrame
    daily: pd.DataFrame
    loaded_at: datetime
    generated_at: datetime
    categories: tuple[dict[str, object], ...]
    category_config_sha256: str | None


def normalize_frame(frame: pd.DataFrame) -> pd.DataFrame:
    normalized = frame.copy()
    if "category" not in normalized:
        normalized["category"] = "all"
    for column in NUMERIC_COLUMNS:
        normalized[column] = pd.to_numeric(
            normalized[column],
            errors="coerce",
        ).fillna(0)
    return normalized


def filter_metrics(
    frame: pd.DataFrame,
    version: str,
    market_type: str,
    category: str = "all",
) -> pd.DataFrame:
    filtered = frame[
        (frame["version"] == version)
        & (frame["market_type"] == market_type)
        & (frame["category"] == category)
    ].copy()
    order_column = "date" if "date" in filtered.columns else "month"
    return filtered.sort_values(order_column).reset_index(drop=True)


def available_market_types(
    monthly: pd.DataFrame,
    version: str,
) -> list[str]:
    order = {"all": 0, "standard": 1, "neg_risk": 2}
    values = monthly.loc[
        monthly["version"] == version,
        "market_type",
    ].drop_duplicates()
    return sorted(values.astype(str).tolist(), key=lambda item: order[item])


def available_months(
    monthly: pd.DataFrame,
    version: str,
    market_type: str,
    category: str = "all",
) -> list[str]:
    filtered = filter_metrics(monthly, version, market_type, category)
    return filtered["month"].astype(str).drop_duplicates().tolist()


def available_categories(
    monthly: pd.DataFrame,
    categories: tuple[dict[str, object], ...],
    version: str,
    market_type: str,
) -> list[str]:
    present = set(
        monthly.loc[
            (monthly["version"] == version)
            & (monthly["market_type"] == market_type),
            "category",
        ].astype(str)
    )
    ordered = [
        str(item["key"])
        for item in sorted(categories, key=lambda item: int(item["order"]))
        if bool(item["enabled"]) and str(item["key"]) in present
    ]
    return ordered or ["all"]


def _fallback_categories() -> tuple[dict[str, object], ...]:
    return (
        {
            "key": "all",
            "order": 0,
            "enabled": True,
            "aggregate_only": True,
            "labels": {"en": "All", "zh": "全部"},
            "tag_ids": [],
        },
    )


def load_categories_for_state(
    state: dict[str, object],
    cache_dir: Path,
) -> tuple[tuple[dict[str, object], ...], str | None]:
    metadata = state.get("categories")
    if not isinstance(metadata, dict):
        return _fallback_categories(), None
    bucket = metadata.get("bucket")
    remote_path = metadata.get("path")
    expected = metadata.get("config_sha256")
    if not all(isinstance(value, str) and value for value in (bucket, remote_path)):
        raise ValueError("aggregation_state 分类配置地址无效")
    if not isinstance(expected, str) or len(expected) != 64:
        raise ValueError("aggregation_state 分类配置 SHA-256 无效")
    destination = cache_dir / f"categories-{expected}.json"
    download_bucket_files(
        str(bucket),
        files=[(str(remote_path), destination)],
    )
    raw = destination.read_bytes()
    actual = hashlib.sha256(raw).hexdigest()
    if actual != expected:
        raise RuntimeError(
            f"分类配置 SHA-256 不匹配：expected={expected}, actual={actual}"
        )
    payload = json.loads(raw)
    categories = payload.get("categories")
    if payload.get("schema_version") != 1 or not isinstance(categories, list):
        raise ValueError("不支持的分类配置 schema")
    normalized = tuple(
        item
        for item in categories
        if isinstance(item, dict)
        and isinstance(item.get("key"), str)
        and isinstance(item.get("order"), int)
        and isinstance(item.get("enabled"), bool)
        and isinstance(item.get("labels"), dict)
    )
    if not normalized or not any(item["key"] == "all" for item in normalized):
        raise ValueError("分类配置缺少 all")
    return normalized, expected


def aggregate_period_metrics(
    monthly: pd.DataFrame,
    period: Period,
) -> pd.DataFrame:
    """Roll additive monthly metrics up to UTC calendar periods."""
    if period not in ("year", "quarter"):
        raise ValueError(f"Unsupported period: {period}")
    if monthly.empty:
        return pd.DataFrame(columns=["period", *NUMERIC_COLUMNS])

    aggregated = monthly.copy()
    months = pd.to_datetime(
        aggregated["month"],
        format="%Y-%m",
        errors="raise",
    )
    years = months.dt.year.astype(str)
    if period == "year":
        aggregated["period"] = years
    else:
        aggregated["period"] = years + "-Q" + months.dt.quarter.astype(str)

    return (
        aggregated.groupby("period", as_index=False, sort=True)[
            list(NUMERIC_COLUMNS)
        ]
        .sum()
        .sort_values("period")
        .reset_index(drop=True)
    )


class AnalyticsStore:
    def __init__(
        self,
        bucket_id: str = BUCKET_ID,
        cache_dir: Path = CACHE_DIR,
    ) -> None:
        self.bucket_id = bucket_id
        self.cache_dir = cache_dir
        self._snapshot: Snapshot | None = None
        self._lock = threading.Lock()

    def refresh(self) -> Snapshot:
        with self._lock:
            self.cache_dir.mkdir(parents=True, exist_ok=True)
            monthly_path = self.cache_dir / MONTHLY_FILE
            daily_path = self.cache_dir / DAILY_FILE
            state_path = self.cache_dir / STATE_FILE
            download_bucket_files(
                self.bucket_id,
                files=[
                    (MONTHLY_FILE, monthly_path),
                    (DAILY_FILE, daily_path),
                    (STATE_FILE, state_path),
                ],
            )
            state = json.loads(state_path.read_text(encoding="utf-8"))
            categories, category_sha = load_categories_for_state(
                state,
                self.cache_dir,
            )
            monthly = normalize_frame(pd.read_parquet(monthly_path))
            daily = normalize_frame(pd.read_parquet(daily_path))
            generated = pd.concat(
                [
                    pd.to_datetime(monthly["generated_at"], utc=True),
                    pd.to_datetime(daily["generated_at"], utc=True),
                ]
            ).max()
            self._snapshot = Snapshot(
                monthly=monthly,
                daily=daily,
                loaded_at=datetime.now(UTC),
                generated_at=generated.to_pydatetime(),
                categories=categories,
                category_config_sha256=category_sha,
            )
            return self._snapshot

    def get(self) -> Snapshot:
        if self._snapshot is None:
            return self.refresh()
        return self._snapshot