#!/usr/bin/env python3
"""Download a structured_analysis parquet and build the local grading template.

Default target is the tiny development run::

    uv run python hf-spaces/grading-answers/ingest_structured_analysis.py

That pulls the latest ``new-llm-judge/tiny-example/en_us/*.parquet`` from
``TransLegal/structured_analysis`` (branch ``dev``) and writes::

    hf-spaces/grading-answers/data/en-us/grading_template.parquet

Pass ``--input`` to skip Hugging Face and build from a local parquet instead.
Use ``--poll`` while a Dagster run is still publishing.
"""

from __future__ import annotations

import argparse
import logging
import sys
import time
from pathlib import Path

import pandas as pd
from huggingface_hub import hf_hub_download

_REPO_ROOT = Path(__file__).resolve().parents[2]
if str(_REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(_REPO_ROOT))

# Sibling module (same directory).
from build_grading_template import (  # noqa: E402
    build_grading_template,
    build_template_metadata,
    log_concept_balance,
    log_shortfall,
    write_template_metadata,
)

from config import get_settings  # noqa: E402
from config.constants import REPOS  # noqa: E402
from utils.huggingface import get_hf_api_wrapper  # noqa: E402

logger = logging.getLogger(__name__)

_SPACE_ROOT = Path(__file__).resolve().parent
_DEFAULT_OUTPUT = _SPACE_ROOT / "data" / "en-us" / "grading_template.parquet"

_STRUCTURED_ANALYSIS_REPO = REPOS["structured_answers"]


def _hyphen_to_underscore(jurisdiction: str) -> str:
    return jurisdiction.replace("-", "_")


def _list_run_parquets(
    *,
    branch: str,
    experiment: str,
    variant: str,
    jurisdiction: str,
) -> list[str]:
    """Return matching parquet paths under experiment/variant/jurisdiction/."""
    prefix = f"{experiment}/{variant}/{_hyphen_to_underscore(jurisdiction)}/"
    hf = get_hf_api_wrapper()
    files = hf.api.list_repo_files(
        repo_id=f"{hf.config.organization}/{_STRUCTURED_ANALYSIS_REPO}",
        repo_type="dataset",
        revision=branch,
    )
    return sorted(f for f in files if f.startswith(prefix) and f.endswith(".parquet"))


def resolve_hf_parquet(
    *,
    branch: str,
    experiment: str,
    variant: str,
    jurisdiction: str,
    path: str | None,
) -> str:
    """Return the repo-relative parquet path to download."""
    if path:
        return path
    matches = _list_run_parquets(
        branch=branch,
        experiment=experiment,
        variant=variant,
        jurisdiction=jurisdiction,
    )
    if not matches:
        raise FileNotFoundError(
            f"No structured_analysis parquet yet under "
            f"{experiment}/{variant}/{_hyphen_to_underscore(jurisdiction)}/ "
            f"on branch {branch!r}. Wait for the pipeline publish, or pass --input."
        )
    chosen = matches[-1]
    if len(matches) > 1:
        logger.info("Found %d parquet(s); using latest: %s", len(matches), chosen)
    return chosen


def download_hf_parquet(
    *,
    repo_path: str,
    branch: str,
    local_dir: Path,
) -> Path:
    settings = get_settings()
    org = settings.huggingface.organization
    token = settings.huggingface.token
    local_dir.mkdir(parents=True, exist_ok=True)
    downloaded = hf_hub_download(
        repo_id=f"{org}/{_STRUCTURED_ANALYSIS_REPO}",
        filename=repo_path,
        repo_type="dataset",
        revision=branch,
        token=token,
        local_dir=str(local_dir),
    )
    return Path(downloaded)


def poll_until_available(
    *,
    branch: str,
    experiment: str,
    variant: str,
    jurisdiction: str,
    path: str | None,
    interval_s: float,
    timeout_s: float,
) -> str:
    """Block until a matching parquet appears on HF (or timeout)."""
    deadline = time.monotonic() + timeout_s
    last_error: Exception | None = None
    while time.monotonic() < deadline:
        try:
            return resolve_hf_parquet(
                branch=branch,
                experiment=experiment,
                variant=variant,
                jurisdiction=jurisdiction,
                path=path,
            )
        except FileNotFoundError as exc:
            last_error = exc
            remaining = max(0, int(deadline - time.monotonic()))
            logger.info(
                "Waiting for HF publish (%ds left)… %s",
                remaining,
                exc,
            )
            time.sleep(interval_s)
    raise TimeoutError(
        f"Timed out after {timeout_s:.0f}s waiting for structured_analysis on "
        f"branch {branch!r}."
    ) from last_error


def build_and_write(
    *,
    input_parquet: Path,
    output_parquet: Path,
    all_fields_of_law: bool,
    max_terms: int | None,
    include_failed: bool,
    terms_per_field: int | None,
    jurisdiction: str,
) -> pd.DataFrame:
    logger.info("Loading %s", input_parquet)
    structured = pd.read_parquet(input_parquet)
    build = build_grading_template(
        structured,
        only_ok=not include_failed,
        balanced=not all_fields_of_law,
        terms_per_field=terms_per_field,
        max_terms=max_terms,
    )
    template = build.template
    output_parquet.parent.mkdir(parents=True, exist_ok=True)
    template.to_parquet(output_parquet, index=False)
    metadata = build_template_metadata(
        input_parquet=input_parquet,
        output_parquet=output_parquet,
        build=build,
        structured_input=structured,
        only_ok=not include_failed,
        jurisdiction=jurisdiction,
        max_terms=max_terms,
    )
    meta_path = write_template_metadata(metadata, output_parquet)
    log_concept_balance(metadata["selection"]["concept_balance"])
    if build.assignment is not None:
        log_shortfall(metadata["selection"]["shortfall"])
    n_results = (
        template[["term", "field_of_law"]].drop_duplicates().shape[0]
        if not template.empty
        else 0
    )
    logger.info(
        "Wrote %d grading rows for %d results → %s",
        len(template),
        n_results,
        output_parquet,
    )
    logger.info("Wrote metadata → %s", meta_path)
    if not template.empty:
        logger.info("Concepts: %s", sorted(template["concept"].unique()))
        logger.info(
            "Questions per concept: %s",
            template.groupby("concept")["question_id"].nunique().to_dict(),
        )
    return template


def _parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description=(
            "Ingest structured_analysis into the local grading app "
            "(download from HF and/or build grading_template.parquet)."
        )
    )
    parser.add_argument(
        "--input",
        type=Path,
        default=None,
        help="Local structured_analysis parquet (skips Hugging Face download).",
    )
    parser.add_argument(
        "--output",
        type=Path,
        default=_DEFAULT_OUTPUT,
        help=f"grading_template.parquet path (default: {_DEFAULT_OUTPUT}).",
    )
    parser.add_argument(
        "--branch",
        default="dev",
        help="HF dataset branch (default: dev).",
    )
    parser.add_argument(
        "--experiment",
        default="new-llm-judge",
        help="Experiment folder on HF (default: new-llm-judge).",
    )
    parser.add_argument(
        "--variant",
        default="tiny-example",
        help="Variant folder on HF (default: tiny-example).",
    )
    parser.add_argument(
        "--jurisdiction",
        default="en-us",
        help="Jurisdiction code with hyphen (default: en-us).",
    )
    parser.add_argument(
        "--hf-path",
        default=None,
        help=(
            "Exact repo-relative parquet path. Default: latest under "
            "experiment/variant/jurisdiction/."
        ),
    )
    parser.add_argument(
        "--download-dir",
        type=Path,
        default=_SPACE_ROOT / "data" / "_downloads",
        help="Where to cache the downloaded HF parquet.",
    )
    parser.add_argument(
        "--poll",
        action="store_true",
        help="Keep checking HF until a matching parquet appears.",
    )
    parser.add_argument(
        "--poll-interval",
        type=float,
        default=30.0,
        help="Seconds between HF checks when --poll is set (default: 30).",
    )
    parser.add_argument(
        "--poll-timeout",
        type=float,
        default=3600.0,
        help="Max seconds to wait with --poll (default: 3600).",
    )
    parser.add_argument(
        "--all-fields-of-law",
        action="store_true",
        help="Keep every (term, field_of_law) row (default: one field per term).",
    )
    parser.add_argument(
        "--max-terms",
        type=int,
        default=None,
        help=(
            "Cap the number of terms graded, taken evenly across concept "
            "classes (default: every term)."
        ),
    )
    parser.add_argument(
        "--include-failed",
        action="store_true",
        help="Include generation_status != ok rows.",
    )
    parser.add_argument(
        "--terms-per-field",
        type=int,
        default=None,
        help="Minimum terms per field of law (default: terms // fields).",
    )
    args = parser.parse_args()
    if args.max_terms is not None and args.max_terms < 0:
        parser.error("--max-terms must not be negative")
    return args


def main() -> None:
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s %(levelname)s %(name)s: %(message)s",
    )
    args = _parse_args()

    if args.input is not None:
        input_parquet = args.input.expanduser().resolve()
        if not input_parquet.is_file():
            raise FileNotFoundError(f"Input parquet not found: {input_parquet}")
    else:
        if args.poll:
            repo_path = poll_until_available(
                branch=args.branch,
                experiment=args.experiment,
                variant=args.variant,
                jurisdiction=args.jurisdiction,
                path=args.hf_path,
                interval_s=args.poll_interval,
                timeout_s=args.poll_timeout,
            )
        else:
            repo_path = resolve_hf_parquet(
                branch=args.branch,
                experiment=args.experiment,
                variant=args.variant,
                jurisdiction=args.jurisdiction,
                path=args.hf_path,
            )
        logger.info(
            "Downloading hf://datasets/.../%s@%s → %s",
            repo_path,
            args.branch,
            args.download_dir,
        )
        input_parquet = download_hf_parquet(
            repo_path=repo_path,
            branch=args.branch,
            local_dir=args.download_dir,
        )

    build_and_write(
        input_parquet=input_parquet,
        output_parquet=args.output.expanduser().resolve(),
        all_fields_of_law=args.all_fields_of_law,
        max_terms=args.max_terms,
        include_failed=args.include_failed,
        terms_per_field=args.terms_per_field,
        jurisdiction=args.jurisdiction,
    )
    logger.info(
        "App-ready. From hf-spaces/grading-answers run: "
        "uv run --with streamlit --with streamlit-shortcuts --with pandas "
        "--with pyarrow --with huggingface-hub --with pycountry --with babel "
        "--with python-dotenv streamlit run src/streamlit_app.py"
    )


if __name__ == "__main__":
    main()