#!/usr/bin/env python3
"""Download HackerOne disclosed reports from HuggingFace and generate a promptfoo dataset.

For each unique program in the sample, imports its real policy from HackerOne
using NoiseGate's importer and saves it to evals/programs/<handle>.yaml.
"""

import argparse
import io
import math
import random
import sys
from pathlib import Path

import httpx
import pandas as pd
import yaml

# Ensure project root is importable
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from noisegate.importer import import_program_from_url
from noisegate.llm import get_default_model
from noisegate.models import Program
from noisegate.program import save_program, load_program

DATASET_ID = "Hacker0x01/hackerone_disclosed_reports"
HF_RESOLVE = f"https://huggingface.co/datasets/{DATASET_ID}/resolve/main/data"

PROGRAMS_DIR = Path(__file__).parent / "programs"

MAX_REPORT_CHARS = 100_000
MIN_REPORT_CHARS = 200

# Substate → expected NoiseGate verdict
SUBSTATE_VERDICT = {
    "resolved": "PASS",
    "not-applicable": "FAIL",
    "informative": "FAIL",
}

# Target proportions by substate (must sum to 1.0)
SUBSTATE_PROPORTIONS = {
    "resolved": 0.6,
    "not-applicable": 0.2,
    "informative": 0.2,
}


def download_parquet(split: str = "train") -> pd.DataFrame:
    """Download parquet file for the given split from HuggingFace."""
    url = f"{HF_RESOLVE}/{split}-00000-of-00001.parquet"
    print(f"Downloading {url} ...")
    resp = httpx.get(url, follow_redirects=True, timeout=120)
    resp.raise_for_status()
    return pd.read_parquet(io.BytesIO(resp.content))


def extract_scope(row) -> tuple[str, str]:
    """Extract asset_identifier and asset_type from structured_scope dict."""
    scope = row.get("structured_scope")
    if isinstance(scope, dict):
        return (
            scope.get("asset_identifier", ""),
            scope.get("asset_type", "URL"),
        )
    return ("", "URL")


def extract_weakness(row) -> str:
    """Extract weakness name from weakness dict."""
    w = row.get("weakness")
    if isinstance(w, dict):
        return w.get("name", "")
    return ""


def extract_program(row) -> str:
    """Extract program handle from team dict."""
    team = row.get("team")
    if isinstance(team, dict):
        return team.get("handle", "") or team.get("name", "")
    return ""


def import_program(handle: str, model: str) -> Program | None:
    """Import a program from HackerOne, caching to evals/programs/<handle>.yaml."""
    program_path = PROGRAMS_DIR / f"{handle}.yaml"
    if program_path.exists():
        print(f"  {handle}: using cached program definition")
        return load_program(program_path)

    url = f"https://hackerone.com/{handle}"
    print(f"  {handle}: importing from {url} ...")
    try:
        program = import_program_from_url(
            seed_url=url,
            model=model,
            on_fetch_start=lambda u: print(f"    fetching {u}"),
            on_warning=lambda w: print(f"    warning: {w}"),
        )
        save_program(program, program_path)
        print(f"  {handle}: saved to {program_path}")
        return program
    except Exception as e:
        print(f"  {handle}: import failed — {e}")
        return None


def build_dataset(df: pd.DataFrame, count: int, seed: int, model: str) -> list[dict]:
    """Filter, sample, import programs, and format rows into promptfoo test cases."""
    rng = random.Random(seed)

    # Extract nested fields
    df = df.copy()
    scope_data = df.apply(extract_scope, axis=1)
    df["asset_identifier"] = scope_data.apply(lambda x: x[0])
    df["asset_type"] = scope_data.apply(lambda x: x[1])
    df["weakness_name"] = df.apply(extract_weakness, axis=1)
    df["program_name"] = df.apply(extract_program, axis=1)

    # Filter: usable reports only
    mask = (
        df["vulnerability_information"].notna()
        & (df["vulnerability_information"].str.len() >= MIN_REPORT_CHARS)
        & df["substate"].isin(SUBSTATE_VERDICT)
        & df["asset_identifier"].astype(bool)
        & df["program_name"].astype(bool)
    )
    df = df[mask].copy()

    print(f"Usable reports after filtering: {len(df)}")
    print(f"  by substate: {df['substate'].value_counts().to_dict()}")

    # Stratified sample by substate
    cases = []
    for substate, proportion in SUBSTATE_PROPORTIONS.items():
        target = math.ceil(count * proportion)
        pool = df[df["substate"] == substate]
        n = min(target, len(pool))
        if n == 0:
            print(f"  WARNING: no rows for substate '{substate}'")
            continue
        sampled = pool.sample(n=n, random_state=rng.randint(0, 2**31))
        cases.extend(sampled.to_dict("records"))
        print(f"  sampled {n} '{substate}' reports")

    rng.shuffle(cases)

    # Import program definitions for each unique program
    print(f"\nImporting program definitions ...")
    PROGRAMS_DIR.mkdir(parents=True, exist_ok=True)
    program_handles = sorted({row["program_name"] for row in cases})
    print(f"  {len(program_handles)} unique programs: {program_handles}")

    imported = set()
    for handle in program_handles:
        program = import_program(handle, model)
        if program is not None:
            imported.add(handle)

    # Keep only cases whose programs imported successfully
    cases = [row for row in cases if row["program_name"] in imported]
    print(f"\n{len(cases)} test cases with valid program definitions")

    # Format as promptfoo test cases
    tests = []
    for row in cases:
        report_text = str(row["vulnerability_information"])[:MAX_REPORT_CHARS]
        substate = row["substate"]
        tests.append({
            "vars": {
                "report_id": int(row["id"]),
                "title": str(row.get("title", "")),
                "report_text": report_text,
                "weakness_name": str(row.get("weakness_name", "")),
                "program_name": str(row["program_name"]),
                "asset_identifier": str(row["asset_identifier"]),
                "asset_type": str(row["asset_type"]),
                "substate": substate,
                "expected_verdict": SUBSTATE_VERDICT[substate],
            },
        })

    return tests


def main():
    parser = argparse.ArgumentParser(description="Generate promptfoo dataset from HackerOne reports")
    parser.add_argument("--count", type=int, default=10, help="Total number of test cases (default: 10)")
    parser.add_argument("--output", type=str, default=str(Path(__file__).parent / "dataset.yaml"),
                        help="Output YAML path (default: evals/dataset.yaml)")
    parser.add_argument("--seed", type=int, default=42, help="Random seed for reproducible sampling")
    parser.add_argument("--split", type=str, default="train", help="Dataset split to use (default: train)")
    parser.add_argument("--model", type=str, default=None, help="Model for program import (default: configured)")
    args = parser.parse_args()

    model = args.model or get_default_model()

    df = download_parquet(split=args.split)
    print(f"Downloaded {len(df)} reports")

    tests = build_dataset(df, count=args.count, seed=args.seed, model=model)
    print(f"\nGenerated {len(tests)} test cases")

    output_path = Path(args.output)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    with open(output_path, "w") as f:
        yaml.dump(tests, f, default_flow_style=False, allow_unicode=True, width=120)

    print(f"Written to {output_path}")


if __name__ == "__main__":
    main()
