#!/usr/bin/env python3
"""
Creates a random subset of metadata for debugging use in training.

This script:
1. Random samples rows from the input meta file
2. For each sampled row, includes any rows that are referenced via cover_ids
3. Outputs the combined subset for debugging

Usage:
    python make_mini_meta.py <input_file> <output_file> [--sample_size N] [--max_rows N]
"""

import orjson
import json
import random
import argparse
from pathlib import Path
from tqdm import tqdm
from typing import Dict, Set, List


def load_jsonl_limited(file_path: str, max_rows: int = None) -> List[Dict]:
    """Load JSONL file with optional row limit."""
    data = []
    with open(file_path, "r") as f:
        for i, line in enumerate(tqdm(f, desc=f"Loading {Path(file_path).name}")):
            if max_rows and i >= max_rows:
                break
            data.append(orjson.loads(line))
    return data


def create_id_to_row_map(data: List[Dict]) -> Dict[str, Dict]:
    """Create mapping from id to row data."""
    return {row["id"]: row for row in data}


def get_referenced_ids(row: Dict, valid_ids: Set[str]) -> Set[str]:
    """Extract all IDs that this row references by checking all keys and values."""
    referenced = set()

    def extract_ids_from_value(value):
        """Recursively extract potential IDs from any value."""
        if isinstance(value, str):
            # Check if this string is a valid ID
            if value in valid_ids:
                referenced.add(value)
        elif isinstance(value, list):
            for item in value:
                extract_ids_from_value(item)
        elif isinstance(value, dict):
            for v in value.values():
                extract_ids_from_value(v)

    # Extract from all values in the row (skip the 'id' key itself)
    for key, value in row.items():
        if key != "id":  # Don't include the row's own ID
            extract_ids_from_value(value)

    return referenced


def collect_all_referenced_rows(sampled_rows: List[Dict], id_to_row: Dict[str, Dict]) -> Set[str]:
    """Collect all rows that are referenced by the sampled rows."""
    all_referenced_ids = set()
    valid_ids = set(id_to_row.keys())

    for row in sampled_rows:
        referenced_ids = get_referenced_ids(row, valid_ids)
        all_referenced_ids.update(referenced_ids)

    return all_referenced_ids


def write_jsonl(data: List[Dict], output_file: str):
    """Write data to JSONL file."""
    with open(output_file, "w") as f:
        for row in data:
            f.write(json.dumps(row) + "\n")


def main():
    parser = argparse.ArgumentParser(description="Create a random subset of metadata for debugging")
    parser.add_argument("input_file", help="Input JSONL file path")
    parser.add_argument("output_file", help="Output JSONL file path")
    parser.add_argument(
        "--sample_size",
        type=int,
        default=100000,
        help="Number of rows to randomly sample (default: 1000)",
    )
    parser.add_argument(
        "--max_rows",
        type=int,
        default=100000,
        help="Maximum rows to load from input file (default: 100000)",
    )
    parser.add_argument(
        "--seed", type=int, default=42, help="Random seed for reproducibility (default: 42)"
    )
    parser.add_argument(
        "--force_include_id", type=str, help="Force include a specific ID in the sample for testing"
    )

    args = parser.parse_args()

    # Set random seed
    random.seed(args.seed)

    print(f"Loading data from {args.input_file} (max {args.max_rows} rows)...")
    data = load_jsonl_limited(args.input_file, args.max_rows)

    print(f"Loaded {len(data)} rows")
    print(f"Creating ID to row mapping...")
    id_to_row = create_id_to_row_map(data)

    print(f"Random sampling {args.sample_size} rows...")
    if args.sample_size >= len(data):
        print(f"Sample size ({args.sample_size}) >= data size ({len(data)}), using all data")
        sampled_rows = data
    else:
        sampled_rows = random.sample(data, args.sample_size)

    # Force include specific ID if requested
    if args.force_include_id:
        if args.force_include_id in id_to_row:
            forced_row = id_to_row[args.force_include_id]
            if forced_row not in sampled_rows:
                sampled_rows.append(forced_row)
            print(f"Force included ID: {args.force_include_id}")
        else:
            print(f"WARNING: Force include ID {args.force_include_id} not found in dataset")

    print(f"Finding all referenced rows...")
    referenced_ids = collect_all_referenced_rows(sampled_rows, id_to_row)

    if referenced_ids:
        print(f"Found {len(referenced_ids)} referenced IDs: {list(referenced_ids)[:10]}...")
    else:
        print("No referenced IDs found - checking a few sampled rows for potential references...")
        valid_ids = set(id_to_row.keys())
        for i, row in enumerate(sampled_rows[:3]):
            refs = get_referenced_ids(row, valid_ids)
            if refs:
                print(f"  Sample row {i} ({row['id']}) references: {refs}")
            else:
                # Show what string values we're checking
                all_strings = []

                def collect_strings(obj):
                    if isinstance(obj, str) and len(obj) == 11:  # ID length
                        all_strings.append(obj)
                    elif isinstance(obj, list):
                        for item in obj:
                            collect_strings(item)
                    elif isinstance(obj, dict):
                        for v in obj.values():
                            collect_strings(v)

                collect_strings(row)
                print(
                    f"  Sample row {i} ({row['id']}) has {len(all_strings)} 11-char strings: {all_strings[:5]}..."
                )

    # Get the referenced rows
    referenced_rows = [id_to_row[ref_id] for ref_id in referenced_ids]

    # Combine sampled and referenced rows, removing duplicates
    sampled_ids = {row["id"] for row in sampled_rows}
    final_data = sampled_rows.copy()

    for ref_row in referenced_rows:
        if ref_row["id"] not in sampled_ids:
            final_data.append(ref_row)

    print(
        f"Final dataset: {len(sampled_rows)} sampled + {len(referenced_rows)} referenced = {len(final_data)} total rows"
    )

    print(f"Writing to {args.output_file}...")
    write_jsonl(final_data, args.output_file)

    print("Done!")


if __name__ == "__main__":
    main()