"""Extract sample records from JSONL files for manual review."""

import json
import random
from pathlib import Path
from typing import Any, Dict, List, Set

from tqdm import tqdm


def extract_samples_from_jsonl(
    jsonl_path: Path,
    target_ids: Set[str],
    sample_size: int = 20,
    random_seed: int = 42,
) -> List[Dict[str, Any]]:
    """Extract sample records from a JSONL file based on target IDs.

    Args:
        jsonl_path: Path to the source JSONL file
        target_ids: Set of IDs to extract
        sample_size: Maximum number of samples to extract
        random_seed: Random seed for reproducible sampling

    Returns:
        List of sample records (dictionaries)
    """
    if not jsonl_path.exists():
        print(f"Warning: JSONL file not found: {jsonl_path}")
        return []

    # Sample IDs if there are more than sample_size
    if len(target_ids) > sample_size:
        random.seed(random_seed)
        sampled_ids = set(random.sample(list(target_ids), sample_size))
    else:
        sampled_ids = target_ids

    samples = []
    found_ids = set()

    print(f"Extracting up to {len(sampled_ids)} samples from {jsonl_path}...")

    try:
        # Get file size for progress bar
        file_size = jsonl_path.stat().st_size

        with open(jsonl_path, "r") as f:
            with tqdm(total=file_size, unit="B", unit_scale=True) as pbar:
                for line in f:
                    pbar.update(len(line.encode("utf-8")))

                    # Stop early if we found all samples
                    if len(found_ids) >= len(sampled_ids):
                        break

                    try:
                        record = json.loads(line.strip())
                        record_id = record.get("id")

                        if record_id in sampled_ids and record_id not in found_ids:
                            samples.append(record)
                            found_ids.add(record_id)

                    except json.JSONDecodeError as e:
                        print(f"Warning: Failed to parse line: {e}")
                        continue

    except Exception as e:
        print(f"Error reading JSONL file {jsonl_path}: {e}")
        return []

    print(f"Found {len(samples)} samples out of {len(sampled_ids)} requested")

    if len(samples) < len(sampled_ids):
        missing_ids = sampled_ids - found_ids
        print(
            f"Warning: {len(missing_ids)} IDs not found in JSONL file. "
            f"Examples: {list(missing_ids)[:5]}"
        )

    return samples


def save_samples_to_jsonl(samples: List[Dict[str, Any]], output_path: Path) -> None:
    """Save sample records to a JSONL file.

    Args:
        samples: List of sample records
        output_path: Path to save the JSONL file
    """
    output_path.parent.mkdir(parents=True, exist_ok=True)

    with open(output_path, "w") as f:
        for sample in samples:
            f.write(json.dumps(sample) + "\n")

    print(f"Saved {len(samples)} samples to {output_path}")


def extract_and_save_samples(
    jsonl_path: Path,
    added_ids: Set[str],
    removed_ids: Set[str],
    output_dir: Path,
    sample_size: int = 20,
) -> Dict[str, int]:
    """Extract and save samples for both added and removed IDs.

    Args:
        jsonl_path: Path to the source JSONL file
        added_ids: Set of IDs that were added
        removed_ids: Set of IDs that were removed
        output_dir: Directory to save sample files
        sample_size: Number of samples per category

    Returns:
        Dictionary with counts of extracted samples
    """
    output_dir.mkdir(parents=True, exist_ok=True)

    results = {"added_samples": 0, "removed_samples": 0}

    # Extract added samples
    if added_ids:
        print(f"\nExtracting {sample_size} added samples...")
        added_samples = extract_samples_from_jsonl(jsonl_path, added_ids, sample_size=sample_size)
        if added_samples:
            save_samples_to_jsonl(added_samples, output_dir / "added_samples.jsonl")
            results["added_samples"] = len(added_samples)
    else:
        print("No added IDs to sample")

    # Extract removed samples
    if removed_ids:
        print(f"\nExtracting {sample_size} removed samples...")
        removed_samples = extract_samples_from_jsonl(jsonl_path, removed_ids, sample_size=sample_size)
        if removed_samples:
            save_samples_to_jsonl(removed_samples, output_dir / "removed_samples.jsonl")
            results["removed_samples"] = len(removed_samples)
    else:
        print("No removed IDs to sample")

    return results


def analyze_sample_characteristics(
    samples: List[Dict[str, Any]], label: str = "samples"
) -> Dict[str, Any]:
    """Analyze characteristics of sample records.

    Args:
        samples: List of sample records
        label: Label for the samples (e.g., "added", "removed")

    Returns:
        Dictionary with sample characteristics
    """
    if not samples:
        return {}

    # Count records with various fields
    has_tags = sum(1 for s in samples if s.get("tags"))
    has_text = sum(1 for s in samples if s.get("text"))
    has_lang = sum(1 for s in samples if s.get("lang"))
    has_stems = sum(1 for s in samples if s.get("stems"))
    has_weight = sum(1 for s in samples if s.get("weight"))

    # Extract unique tags
    all_tags = set()
    for sample in samples:
        tags = sample.get("tags", [])
        if tags:
            all_tags.update(tags)

    # Extract unique languages
    languages = set(s.get("lang") for s in samples if s.get("lang"))

    # Extract source paths
    local_paths = set()
    s3_paths = set()
    for sample in samples:
        local_path = sample.get("local_filepath", "")
        s3_path = sample.get("s3_filepath", "")

        if local_path:
            # Extract directory portion
            local_dir = str(Path(local_path).parent.parent) if local_path else ""
            if local_dir:
                local_paths.add(local_dir)

        if s3_path:
            # Extract bucket/dataset portion
            if s3_path.startswith("s3://"):
                parts = s3_path[5:].split("/")
                if len(parts) >= 3:
                    s3_prefix = f"s3://{parts[0]}/{parts[1]}/{parts[2]}"
                    s3_paths.add(s3_prefix)

    # Extract weight distribution
    weights = [s.get("weight") for s in samples if s.get("weight") is not None]
    weight_dist = {}
    for w in weights:
        weight_dist[w] = weight_dist.get(w, 0) + 1

    return {
        "label": label,
        "count": len(samples),
        "records_with_tags": has_tags,
        "records_with_text": has_text,
        "records_with_lang": has_lang,
        "records_with_stems": has_stems,
        "records_with_weight": has_weight,
        "unique_tags_count": len(all_tags),
        "sample_tags": list(all_tags)[:20],
        "unique_languages": list(languages),
        "local_path_patterns": list(local_paths),
        "s3_path_patterns": list(s3_paths),
        "weight_distribution": weight_dist,
    }


def create_sample_analysis_report(
    added_samples: List[Dict[str, Any]],
    removed_samples: List[Dict[str, Any]],
    output_dir: Path,
) -> None:
    """Create analysis report for sample records.

    Args:
        added_samples: List of added sample records
        removed_samples: List of removed sample records
        output_dir: Directory to save the report
    """
    added_analysis = analyze_sample_characteristics(added_samples, "added")
    removed_analysis = analyze_sample_characteristics(removed_samples, "removed")

    report = {
        "added_samples": added_analysis,
        "removed_samples": removed_analysis,
    }

    report_path = output_dir / "sample_analysis.json"
    with open(report_path, "w") as f:
        json.dump(report, f, indent=2)

    print(f"Sample analysis report saved to {report_path}")


def extract_full_samples_workflow(
    jsonl_path: Path,
    added_ids: Set[str],
    removed_ids: Set[str],
    output_dir: Path,
    sample_size: int = 20,
) -> None:
    """Full workflow for extracting samples and creating analysis.

    Args:
        jsonl_path: Path to the source JSONL file
        added_ids: Set of IDs that were added
        removed_ids: Set of IDs that were removed
        output_dir: Directory to save all outputs
        sample_size: Number of samples per category
    """
    # Extract samples
    print("\n" + "=" * 60)
    print("EXTRACTING SAMPLES")
    print("=" * 60)

    added_samples = []
    removed_samples = []

    if added_ids:
        print(f"\nExtracting {sample_size} added samples...")
        added_samples = extract_samples_from_jsonl(jsonl_path, added_ids, sample_size=sample_size)
        if added_samples:
            save_samples_to_jsonl(added_samples, output_dir / "added_samples.jsonl")

    if removed_ids:
        print(f"\nExtracting {sample_size} removed samples...")
        removed_samples = extract_samples_from_jsonl(jsonl_path, removed_ids, sample_size=sample_size)
        if removed_samples:
            save_samples_to_jsonl(removed_samples, output_dir / "removed_samples.jsonl")

    # Create analysis report
    if added_samples or removed_samples:
        print("\nCreating sample analysis report...")
        create_sample_analysis_report(added_samples, removed_samples, output_dir)