fisheries_hf.py

"""Restricted Hugging Face fisheries file reader and CSV export helpers.

The module deliberately exposes no arbitrary URL, shell, or filesystem access.
Callers must first validate a path against the configured Dataset live tree and
pass the expected file size to :func:`download_dataset_file`.
"""
from __future__ import annotations

import csv
import hashlib
import io
import itertools
import json
import os
import re
import secrets
import shutil
import time
import zipfile
from collections import defaultdict
from pathlib import Path
from typing import Any, Iterable
from urllib.parse import quote

import httpx


HF_SQUID_DATASET_REPO = (
    os.environ.get("HF_SQUID_DATASET_REPO")
    or os.environ.get("HF_DATASET_REPO")
    or "globalsquiddatabase/squid_dataset"
).strip()
HF_TUNA_DATASET_REPO = (
    os.environ.get("HF_TUNA_DATASET_REPO")
    or "globalsquiddatabase/Tuna-Fisheries-Dataset"
).strip()
HF_DATASET_REPOS = {
    "squid": HF_SQUID_DATASET_REPO,
    "tuna": HF_TUNA_DATASET_REPO,
}
# Backwards-compatible default used when old callers omit a repository.
HF_DATASET_REPO = HF_SQUID_DATASET_REPO
LOCAL_SQUID_DATA_ROOT = Path(os.environ.get("LOCAL_SQUID_DATA_ROOT", "").strip()).expanduser() if os.environ.get("LOCAL_SQUID_DATA_ROOT", "").strip() else None
LOCAL_TUNA_DATA_ROOT = Path(os.environ.get("LOCAL_TUNA_DATA_ROOT", "").strip()).expanduser() if os.environ.get("LOCAL_TUNA_DATA_ROOT", "").strip() else None
LOCAL_DATA_ROOTS = {"squid": LOCAL_SQUID_DATA_ROOT, "tuna": LOCAL_TUNA_DATA_ROOT}
HF_SQUID_DATASET_REVISION = (
    os.environ.get("HF_SQUID_DATASET_REVISION")
    or os.environ.get("HF_DATASET_REVISION")
    or ""
).strip()
HF_TUNA_DATASET_REVISION = os.environ.get("HF_TUNA_DATASET_REVISION", "").strip()
HF_TOKEN = os.environ.get("HF_TOKEN", "").strip()
PUBLIC_BASE_URL = os.environ.get("PUBLIC_BASE_URL", "").strip().rstrip("/")

CACHE_ROOT = Path(
    os.environ.get("HF_FISHERIES_CACHE_ROOT", "/tmp/squid_hf_fisheries_cache")
)
EXPORT_ROOT = Path(
    os.environ.get("FISHERIES_EXPORT_ROOT", "/tmp/squid_fisheries_exports")
)
CACHE_ROOT.mkdir(parents=True, exist_ok=True)
EXPORT_ROOT.mkdir(parents=True, exist_ok=True)

MAX_CSV_BYTES = int(os.environ.get("HF_FISHERIES_MAX_CSV_BYTES", 100 * 1024 * 1024))
MAX_ZIP_BYTES = int(os.environ.get("HF_FISHERIES_MAX_ZIP_BYTES", 1536 * 1024 * 1024))
EXPORT_TTL_SECONDS = int(os.environ.get("FISHERIES_EXPORT_TTL_SECONDS", "86400"))

_REVISION_CACHE: dict[str, dict[str, Any]] = {}

YEAR_ALIASES = ("year", "yearc", "年份", "yr")
MONTH_ALIASES = ("month", "月份", "mon")
DATE_ALIASES = ("date", "time", "datetime", "日期", "时间", "year_month")
LON_ALIASES = ("lon", "longitude", "decimal_longitude", "经度", "x")
LAT_ALIASES = ("lat", "latitude", "decimal_latitude", "纬度", "y")


def _headers() -> dict[str, str]:
    return {"Authorization": f"Bearer {HF_TOKEN}"} if HF_TOKEN else {}


def normalize_repository(repository: str | None = None) -> str:
    """Resolve a safe repository selector to one of the two configured repos."""
    value = str(repository or "squid").strip()
    lowered = value.lower()
    if lowered in HF_DATASET_REPOS:
        return HF_DATASET_REPOS[lowered]
    for repo in HF_DATASET_REPOS.values():
        if value == repo:
            return repo
    raise ValueError("repository must be squid, tuna, or an exact configured repository id")


def local_root_for_repository(repository: str | None = None) -> Path | None:
    """Return the configured local mirror root for a logical repository."""
    value = str(repository or "squid").strip()
    domain = value.lower() if value.lower() in LOCAL_DATA_ROOTS else None
    if domain is None:
        for key, repo in HF_DATASET_REPOS.items():
            if value == repo:
                domain = key
                break
    root = LOCAL_DATA_ROOTS.get(domain or "")
    if root and root.is_dir():
        return root.resolve()
    return None


def local_dataset_tree(repository: str | None = None) -> list[dict[str, Any]]:
    """List supported files from a local school-server mirror.

    The mirror is a read-through cache, not an alternative dataset.  Callers
    merge this list with the Hugging Face live tree and prefer these entries
    only when the same path exists in both places.
    """
    root = local_root_for_repository(repository)
    if root is None:
        return []
    items: list[dict[str, Any]] = []
    for path in root.rglob("*"):
        if not path.is_file() or path.suffix.lower() not in {".csv", ".tsv", ".zip"}:
            continue
        items.append({
            "path": path.relative_to(root).as_posix(),
            "type": "file",
            "size": path.stat().st_size,
            "storage_origin": "school_local_mirror",
        })
    return sorted(items, key=lambda item: str(item["path"]))


def merge_local_with_remote_tree(
    remote_items: Iterable[dict[str, Any]], repository: str | None = None
) -> list[dict[str, Any]]:
    """Combine the authoritative HF tree with an optional local read cache.

    Every remote file remains visible.  A local file with the *same* safe
    relative path replaces the remote entry so content reads take the local
    route first; files absent locally keep their Hugging Face origin and can
    be downloaded on demand.  This gives the school and HF deployments the
    same inventory instead of hiding remote-only files on the school server.
    """
    merged: dict[str, dict[str, Any]] = {}
    for item in remote_items:
        if not isinstance(item, dict):
            continue
        path = str(item.get("path") or "").strip().lstrip("/")
        if not path:
            continue
        row = dict(item)
        row["path"] = path
        row.setdefault("storage_origin", "huggingface")
        merged[path] = row
    for item in local_dataset_tree(repository):
        path = str(item.get("path") or "").strip()
        if path:
            merged[path] = dict(item)
    return [merged[path] for path in sorted(merged)]


def _revision(repository: str) -> str:
    pinned = (
        HF_TUNA_DATASET_REVISION
        if repository == HF_TUNA_DATASET_REPO
        else HF_SQUID_DATASET_REVISION
    )
    if pinned:
        return pinned
    now = time.time()
    cache = _REVISION_CACHE.get(repository) or {}
    cached = str(cache.get("sha") or "")
    if cached and now - float(cache.get("ts") or 0) < 300:
        return cached
    url = f"https://huggingface.co/api/datasets/{repository}"
    with httpx.Client(timeout=30.0, follow_redirects=True) as client:
        response = client.get(url, headers=_headers())
    if response.status_code in {401, 403}:
        raise RuntimeError("Hugging Face Dataset 无读取权限，请检查 HF_TOKEN。")
    response.raise_for_status()
    body = response.json()
    sha = str(body.get("sha") or "").strip()
    if not re.fullmatch(r"[0-9a-fA-F]{40}", sha):
        raise RuntimeError("无法获得 Hugging Face Dataset 的固定提交 SHA。")
    _REVISION_CACHE[repository] = {"ts": now, "sha": sha}
    return sha


def _safe_cache_path(repository: str, path: str, revision: str) -> Path:
    suffix = Path(path).suffix.lower()
    digest = hashlib.sha256(f"{repository}:{revision}:{path}".encode()).hexdigest()
    return CACHE_ROOT / f"{digest}{suffix}"


def download_dataset_file(
    path: str,
    expected_size: int,
    repository: str | None = None,
) -> tuple[Path, str]:
    """Download one validated Dataset file to a revision-keyed local cache."""
    clean = str(path or "").strip().lstrip("/")
    if not clean or "\x00" in clean or any(part in {"", ".", ".."} for part in clean.split("/")):
        raise ValueError("invalid dataset path")
    suffix = Path(clean).suffix.lower()
    limit = MAX_ZIP_BYTES if suffix == ".zip" else MAX_CSV_BYTES
    if suffix not in {".csv", ".tsv", ".zip"}:
        raise ValueError("当前内容查询仅支持 CSV、TSV 和包含 CSV/TSV 的 ZIP。")
    if expected_size <= 0:
        raise ValueError("live tree did not provide a positive file size")
    if expected_size > limit:
        raise ValueError(
            f"文件大小 {expected_size} bytes 超过在线处理上限 {limit} bytes；"
            "请先在学校服务器生成查询就绪的分区文件。"
        )

    repository = normalize_repository(repository)
    local_root = local_root_for_repository(repository)
    if local_root is not None:
        candidate = (local_root / clean).resolve()
        if local_root not in candidate.parents and candidate != local_root:
            raise ValueError("invalid local dataset path")
        if candidate.is_file():
            actual_size = candidate.stat().st_size
            if expected_size and actual_size != expected_size:
                raise RuntimeError(f"本地文件大小不一致：expected={expected_size}, actual={actual_size}")
            return candidate, "local"
        # A partial local mirror must never make HF-only files unreachable.
        # Continue to the revision-pinned remote download below.
    revision = _revision(repository)
    target = _safe_cache_path(repository, clean, revision)
    if target.exists() and target.stat().st_size == expected_size:
        return target, revision

    partial = target.with_suffix(target.suffix + ".part")
    partial.unlink(missing_ok=True)
    encoded_path = quote(clean, safe="/")
    url = (
        f"https://huggingface.co/datasets/{repository}/resolve/"
        f"{revision}/{encoded_path}"
    )
    total = 0
    try:
        with httpx.stream(
            "GET",
            url,
            headers=_headers(),
            follow_redirects=True,
            timeout=httpx.Timeout(connect=20, read=300, write=30, pool=30),
        ) as response:
            if response.status_code in {401, 403}:
                raise RuntimeError("Hugging Face 文件无读取权限，请检查 HF_TOKEN。")
            response.raise_for_status()
            with partial.open("wb") as stream:
                for chunk in response.iter_bytes(1024 * 1024):
                    if not chunk:
                        continue
                    total += len(chunk)
                    if total > limit:
                        raise ValueError("download exceeded configured size limit")
                    stream.write(chunk)
        if total != expected_size:
            raise RuntimeError(
                f"文件下载不完整：expected={expected_size}, received={total}"
            )
        partial.replace(target)
    except Exception:
        partial.unlink(missing_ok=True)
        raise
    return target, revision


def _decode_text(path: Path) -> tuple[str, str]:
    raw = path.read_bytes()
    for encoding in ("utf-8-sig", "utf-8", "gb18030"):
        try:
            return raw.decode(encoding), encoding
        except UnicodeDecodeError:
            pass
    return raw.decode("utf-8", errors="replace"), "utf-8-replace"


def _dialect(text: str, suffix: str = ".csv") -> str:
    sample = text[:20000]
    try:
        return csv.Sniffer().sniff(sample, delimiters=",\t;|").delimiter
    except Exception:
        return "\t" if suffix == ".tsv" else ","


def _column(columns: Iterable[str], aliases: Iterable[str]) -> str | None:
    exact = {str(col).strip().lower(): str(col) for col in columns}
    for alias in aliases:
        if alias.lower() in exact:
            return exact[alias.lower()]
    return None


def _number(value: Any) -> float | None:
    text = str(value or "").strip().replace(",", "")
    if not text:
        return None
    try:
        return float(text)
    except Exception:
        return None


def _year_month(row: dict[str, str], columns: list[str]) -> tuple[int | None, int | None]:
    year_col = _column(columns, YEAR_ALIASES)
    month_col = _column(columns, MONTH_ALIASES)
    date_col = _column(columns, DATE_ALIASES)
    year = None
    month = None
    if year_col:
        value = _number(row.get(year_col))
        if value is not None and 1800 <= int(value) <= 2200:
            year = int(value)
    if month_col:
        value = _number(row.get(month_col))
        if value is not None and 1 <= int(value) <= 12:
            month = int(value)
    if date_col and (year is None or month is None):
        text = str(row.get(date_col) or "")
        match = re.search(r"(19\d{2}|20\d{2}|21\d{2})[-/]?(0?[1-9]|1[0-2])?", text)
        if match:
            year = year or int(match.group(1))
            month = month or (int(match.group(2)) if match.group(2) else None)
    return year, month


def _metric_columns(columns: list[str], requested: str | None = None) -> list[str]:
    if requested:
        wanted = [x.strip() for x in requested.split(",") if x.strip()]
        missing = [x for x in wanted if x not in columns]
        if missing:
            raise ValueError(f"requested metric columns not found: {missing}")
        return wanted
    keys = (
        "effort", "fishing_hour", "fishing hours", "apparent_fishing",
        "catch", "harvest", "landing", "cpue", "value",
    )
    excluded = set(YEAR_ALIASES + MONTH_ALIASES + LON_ALIASES + LAT_ALIASES)
    return [
        col for col in columns
        if str(col).strip().lower() not in excluded
        and any(key in str(col).strip().lower() for key in keys)
    ]


def _matches_filters(
    row: dict[str, str],
    columns: list[str],
    *,
    year: int | None,
    lon_min: float | None,
    lon_max: float | None,
    lat_min: float | None,
    lat_max: float | None,
) -> tuple[bool, int | None, int | None]:
    row_year, row_month = _year_month(row, columns)
    if year is not None and row_year != int(year):
        return False, row_year, row_month
    lon_col = _column(columns, LON_ALIASES)
    lat_col = _column(columns, LAT_ALIASES)
    if any(v is not None for v in (lon_min, lon_max)):
        if not lon_col:
            raise ValueError("经度筛选已请求，但文件中未识别到经度字段。")
        lon = _number(row.get(lon_col))
        if lon is None or (lon_min is not None and lon < lon_min) or (lon_max is not None and lon > lon_max):
            return False, row_year, row_month
    if any(v is not None for v in (lat_min, lat_max)):
        if not lat_col:
            raise ValueError("纬度筛选已请求，但文件中未识别到纬度字段。")
        lat = _number(row.get(lat_col))
        if lat is None or (lat_min is not None and lat < lat_min) or (lat_max is not None and lat > lat_max):
            return False, row_year, row_month
    return True, row_year, row_month


def _create_export(filename: str, rows: Iterable[dict[str, Any]], columns: list[str]) -> dict[str, str]:
    now = time.time()
    for item in list(EXPORT_ROOT.iterdir())[:2000]:
        if not item.is_dir():
            continue
        try:
            meta = json.loads((item / "meta.json").read_text(encoding="utf-8"))
            expired = float(meta.get("expires_ts") or 0) < now
        except Exception:
            expired = True
        if expired:
            shutil.rmtree(item, ignore_errors=True)

    token = secrets.token_urlsafe(24)
    safe_name = re.sub(r"[^A-Za-z0-9._-]+", "_", filename).strip("._") or "fisheries_export.csv"
    folder = EXPORT_ROOT / token
    folder.mkdir(parents=True, exist_ok=False)
    target = folder / safe_name
    with target.open("w", encoding="utf-8-sig", newline="") as stream:
        writer = csv.DictWriter(stream, fieldnames=columns, extrasaction="ignore")
        writer.writeheader()
        writer.writerows(rows)
    meta = {
        "filename": safe_name,
        "content_type": "text/csv; charset=utf-8",
        "size_bytes": target.stat().st_size,
        "created_ts": now,
        "expires_ts": now + EXPORT_TTL_SECONDS,
    }
    (folder / "meta.json").write_text(json.dumps(meta, ensure_ascii=False, indent=2), encoding="utf-8")
    path = f"/api/fisheries/download/{token}"
    return {
        "token": token,
        "filename": safe_name,
        "download_path": path,
        "download_url": f"{PUBLIC_BASE_URL}{path}" if PUBLIC_BASE_URL else path,
    }


def _process_reader(
    reader: csv.DictReader,
    *,
    source_label: str,
    year: int | None,
    lon_min: float | None,
    lon_max: float | None,
    lat_min: float | None,
    lat_max: float | None,
    metric_columns: str | None,
    max_rows: int,
    state: dict[str, Any],
) -> None:
    columns = [str(x or "").strip() for x in (reader.fieldnames or [])]
    if not columns:
        return
    if not state["columns"]:
        state["columns"] = columns
        state["metrics"] = _metric_columns(columns, metric_columns)
    elif columns != state["columns"]:
        state["schema_variants"].append({"source": source_label, "columns": columns})

    seen = state["seen"]
    for raw in reader:
        if state["scanned"] >= max_rows:
            state["truncated"] = True
            return
        state["scanned"] += 1
        row = {str(k or "").strip(): "" if v is None else str(v).strip() for k, v in raw.items()}
        matched, row_year, row_month = _matches_filters(
            row,
            columns,
            year=year,
            lon_min=lon_min,
            lon_max=lon_max,
            lat_min=lat_min,
            lat_max=lat_max,
        )
        if not matched:
            continue
        # Quality counts are reported for the requested result set.  This
        # keeps large year/bbox ZIP queries bounded instead of retaining a
        # duplicate key for every global row that was scanned and rejected.
        key = tuple(row.get(col, "") for col in columns)
        if key in seen:
            state["duplicates"] += 1
        else:
            seen.add(key)
        for col in columns:
            if not row.get(col, "").strip():
                state["missing"][col] += 1
        state["matched"] += 1
        state["rows"].append(row)
        if row_year is not None:
            state["years"].add(row_year)
        if row_month is not None:
            state["months"][row_month] += 1
        if row_year is not None:
            state["annual_counts"][row_year] += 1
            for col in state["metrics"]:
                value = _number(row.get(col))
                if value is not None:
                    state["annual_values"][(row_year, col)] += value


def analyze_and_export(
    local_path: Path,
    *,
    dataset_path: str,
    revision: str,
    repository: str | None = None,
    year: int | None = None,
    lon_min: float | None = None,
    lon_max: float | None = None,
    lat_min: float | None = None,
    lat_max: float | None = None,
    metric_columns: str | None = None,
    max_rows: int = 2_000_000,
) -> dict[str, Any]:
    if lon_min is not None and not -180 <= float(lon_min) <= 180:
        raise ValueError("lon_min 必须位于 -180 至 180。")
    if lon_max is not None and not -180 <= float(lon_max) <= 180:
        raise ValueError("lon_max 必须位于 -180 至 180。")
    if lat_min is not None and not -90 <= float(lat_min) <= 90:
        raise ValueError("lat_min 必须位于 -90 至 90。")
    if lat_max is not None and not -90 <= float(lat_max) <= 90:
        raise ValueError("lat_max 必须位于 -90 至 90。")
    if lon_min is not None and lon_max is not None and float(lon_min) > float(lon_max):
        raise ValueError("lon_min 不能大于 lon_max。")
    if lat_min is not None and lat_max is not None and float(lat_min) > float(lat_max):
        raise ValueError("lat_min 不能大于 lat_max。")

    has_filter = year is not None or any(
        value is not None for value in (lon_min, lon_max, lat_min, lat_max)
    )
    if (
        local_path.suffix.lower() == ".zip"
        and local_path.stat().st_size > 200 * 1024 * 1024
        and not has_filter
    ):
        raise ValueError(
            "大型 ZIP 查询必须提供 year 或经纬度范围，避免无边界解压扫描；"
            "请补充筛选条件后重试。"
        )

    max_rows = max(1, min(int(max_rows or 2_000_000), 5_000_000))
    state: dict[str, Any] = {
        "columns": [], "metrics": [], "schema_variants": [], "scanned": 0,
        "matched": 0, "duplicates": 0, "missing": defaultdict(int),
        "years": set(), "months": defaultdict(int), "annual_counts": defaultdict(int),
        "annual_values": defaultdict(float), "rows": [], "seen": set(), "truncated": False,
    }

    suffix = local_path.suffix.lower()
    members: list[str] = []
    if suffix in {".csv", ".tsv"}:
        text, encoding = _decode_text(local_path)
        delimiter = _dialect(text, suffix)
        _process_reader(
            csv.DictReader(io.StringIO(text), delimiter=delimiter),
            source_label=dataset_path,
            year=year, lon_min=lon_min, lon_max=lon_max,
            lat_min=lat_min, lat_max=lat_max,
            metric_columns=metric_columns, max_rows=max_rows, state=state,
        )
    else:
        encoding = "utf-8-replace"
        delimiter = "auto-by-member-extension"
        with zipfile.ZipFile(local_path) as archive:
            candidates = [
                info for info in archive.infolist()
                if not info.is_dir() and Path(info.filename).suffix.lower() in {".csv", ".tsv"}
            ]
            if year is not None:
                preferred = [info for info in candidates if str(year) in info.filename]
                if preferred:
                    candidates = preferred
            for info in candidates:
                members.append(info.filename)
                with archive.open(info) as binary:
                    text_stream = io.TextIOWrapper(binary, encoding="utf-8-sig", errors="replace", newline="")
                    sample_lines = list(itertools.islice(text_stream, 50))
                    member_delimiter = _dialect(
                        "".join(sample_lines),
                        Path(info.filename).suffix.lower(),
                    )
                    _process_reader(
                        csv.DictReader(
                            itertools.chain(sample_lines, text_stream),
                            delimiter=member_delimiter,
                        ),
                        source_label=info.filename,
                        year=year, lon_min=lon_min, lon_max=lon_max,
                        lat_min=lat_min, lat_max=lat_max,
                        metric_columns=metric_columns, max_rows=max_rows, state=state,
                    )
                if state["truncated"]:
                    break

    if not state["columns"]:
        raise ValueError("未在文件中发现可读取的 CSV/TSV 表格。")

    annual_summary = []
    for yr in sorted(state["annual_counts"]):
        item: dict[str, Any] = {"year": yr, "record_count": state["annual_counts"][yr]}
        for col in state["metrics"]:
            item[f"sum_{col}"] = state["annual_values"].get((yr, col), 0.0)
        annual_summary.append(item)

    # Return three distinct, truthfully labelled artifacts when the source has
    # a usable year field. Keep the raw-detail export as the backwards-
    # compatible primary download_url/token/filename.
    stamp = int(time.time())
    source_stem = Path(dataset_path).stem or "fisheries_query"
    raw_export = _create_export(
        f"{source_stem}_{stamp}_filtered_raw.csv",
        state["rows"],
        state["columns"],
    )
    deduplicated_rows = []
    exported_keys = set()
    for row in state["rows"]:
        key = tuple(row.get(col, "") for col in state["columns"])
        if key in exported_keys:
            continue
        exported_keys.add(key)
        deduplicated_rows.append(row)
    deduplicated_export = _create_export(
        f"{source_stem}_{stamp}_deduplicated.csv",
        deduplicated_rows,
        state["columns"],
    )
    exports = [
        {"kind": "filtered_raw", **raw_export},
        {
            "kind": "deduplicated",
            "record_count": len(deduplicated_rows),
            **deduplicated_export,
        },
    ]
    if annual_summary:
        annual_columns = list(annual_summary[0])
        annual_export = _create_export(
            f"{source_stem}_{stamp}_annual_summary.csv",
            annual_summary,
            annual_columns,
        )
        exports.append({"kind": "annual_summary", **annual_export})

    return {
        "status": "ok",
        "repository": normalize_repository(repository),
        "revision": revision,
        "dataset_path": dataset_path,
        "encoding": encoding,
        "delimiter": delimiter,
        "columns": state["columns"],
        "metric_columns": state["metrics"],
        "scanned_row_count": state["scanned"],
        "matched_row_count": state["matched"],
        "scan_truncated": state["truncated"],
        "time_range": {
            "min_year": min(state["years"]) if state["years"] else None,
            "max_year": max(state["years"]) if state["years"] else None,
        },
        "monthly_record_counts": [
            {"month": month, "record_count": state["months"][month]}
            for month in sorted(state["months"])
        ],
        "missing_values_by_column": {
            col: int(state["missing"].get(col, 0)) for col in state["columns"]
        },
        "exact_duplicate_count": state["duplicates"],
        "deduplicated_record_count": len(deduplicated_rows),
        "annual_summary": annual_summary[:100],
        "exports": exports,
        "download_urls": [item["download_url"] for item in exports],
        "zip_members_processed": members[:200],
        "schema_variants": state["schema_variants"][:20],
        **raw_export,
    }