#!/usr/bin/env python3
"""Main orchestrator script for running all data version comparisons."""

import argparse
import json
import sys
from pathlib import Path
from typing import List, Tuple

# Handle both direct execution and module import
if __name__ == "__main__" and __package__ is None:
    # Add src directory to path for direct execution
    sys.path.insert(0, str(Path(__file__).resolve().parent.parent.parent))

try:
    from sunodata.data_diff.compare_versions import (
        compare_versions,
        save_comparison_results,
    )
    from sunodata.data_diff.extract_samples import extract_full_samples_workflow
    from sunodata.data_diff.load_profiles import (
        get_source_jsonl_path,
        load_version_data,
    )
    from sunodata.data_diff.visualize_diffs import (
        create_all_plots,
        plot_overall_trends,
    )
except ImportError:
    from compare_versions import compare_versions, save_comparison_results
    from extract_samples import extract_full_samples_workflow
    from load_profiles import get_source_jsonl_path, load_version_data
    from visualize_diffs import create_all_plots, plot_overall_trends


def get_version_pairs(versions: List[str]) -> List[Tuple[str, str]]:
    """Generate sequential version pairs.

    Args:
        versions: List of version strings (e.g., ["v0", "v1", "v2", ...])

    Returns:
        List of tuples representing version pairs (e.g., [("v0", "v1"), ("v1", "v2"), ...])
    """
    pairs = []
    for i in range(len(versions) - 1):
        pairs.append((versions[i], versions[i + 1]))
    return pairs


def compare_version_pair(
    profile_dir: Path,
    old_version: str,
    new_version: str,
    output_dir: Path,
    sample_size: int = 20,
    generate_plots: bool = True,
) -> dict:
    """Compare a single version pair and generate all outputs.

    Args:
        profile_dir: Base directory containing all profiles
        old_version: Old version string (e.g., "v0")
        new_version: New version string (e.g., "v1")
        output_dir: Base output directory
        sample_size: Number of samples to extract per category
        generate_plots: Whether to generate visualization plots

    Returns:
        Comparison results dictionary
    """
    print("\n" + "=" * 80)
    print(f"COMPARING {old_version} → {new_version}")
    print("=" * 80)

    # Create output directory for this comparison
    pair_output_dir = output_dir / f"diff_{old_version}_to_{new_version}"
    pair_output_dir.mkdir(parents=True, exist_ok=True)

    # Load data for both versions
    print(f"\nLoading {old_version} data...")
    old_data = load_version_data(profile_dir, old_version, split="train")
    if old_data is None:
        print(f"ERROR: Failed to load {old_version} data")
        return None

    old_analysis, old_ids, old_run_dir = old_data
    print(f"Loaded {old_version}: {len(old_ids):,} IDs from {old_run_dir.parent.name}")

    print(f"\nLoading {new_version} data...")
    new_data = load_version_data(profile_dir, new_version, split="train")
    if new_data is None:
        print(f"ERROR: Failed to load {new_version} data")
        return None

    new_analysis, new_ids, new_run_dir = new_data
    print(f"Loaded {new_version}: {len(new_ids):,} IDs from {new_run_dir.parent.name}")

    # Perform comparison
    print("\nComparing versions...")
    comparison = compare_versions(old_analysis, new_analysis, old_ids, new_ids, old_version, new_version)

    # Save comparison results
    print("\nSaving comparison results...")
    save_comparison_results(comparison, pair_output_dir)

    # Extract samples
    added_ids = set(comparison["id_changes"]["added_ids"])
    removed_ids = set(comparison["id_changes"]["removed_ids"])

    if added_ids or removed_ids:
        # Use the newer version's JSONL for added samples, older for removed
        new_jsonl_path = get_source_jsonl_path(new_analysis)
        old_jsonl_path = get_source_jsonl_path(old_analysis)

        if new_jsonl_path and new_jsonl_path.exists():
            print(f"\nExtracting added samples from {new_jsonl_path}...")
            extract_full_samples_workflow(
                new_jsonl_path,
                added_ids,
                set(),  # No removed samples from new version
                pair_output_dir,
                sample_size=sample_size,
            )

        if old_jsonl_path and old_jsonl_path.exists():
            print(f"\nExtracting removed samples from {old_jsonl_path}...")
            extract_full_samples_workflow(
                old_jsonl_path,
                set(),  # No added samples from old version
                removed_ids,
                pair_output_dir,
                sample_size=sample_size,
            )
    else:
        print("\nNo ID changes detected, skipping sample extraction")

    # Generate plots
    if generate_plots:
        create_all_plots(comparison, pair_output_dir)

    # Print summary
    print("\n" + "-" * 80)
    print("SUMMARY")
    print("-" * 80)
    print(
        f"Total records: {comparison['summary']['old_total_records']:,} → {comparison['summary']['new_total_records']:,}"
    )
    print(f"Record change: {comparison['summary']['record_change']:+,}")
    print(f"Added IDs: {comparison['summary']['added_ids_count']:,}")
    print(f"Removed IDs: {comparison['summary']['removed_ids_count']:,}")
    print(f"Net change: {comparison['summary']['net_change']:+,}")
    print("-" * 80)

    return comparison


def run_all_diffs(
    profile_dir: Path,
    output_dir: Path,
    versions: List[str],
    sample_size: int = 20,
    generate_plots: bool = True,
) -> None:
    """Run all sequential version comparisons.

    Args:
        profile_dir: Base directory containing all profiles
        output_dir: Base output directory
        versions: List of versions to compare sequentially
        sample_size: Number of samples to extract per category
        generate_plots: Whether to generate visualization plots
    """
    output_dir.mkdir(parents=True, exist_ok=True)

    # Get version pairs
    pairs = get_version_pairs(versions)
    print(f"\nWill compare {len(pairs)} version pairs:")
    for old, new in pairs:
        print(f"  {old} → {new}")

    # Run comparisons
    all_comparisons = []
    for old_version, new_version in pairs:
        comparison = compare_version_pair(
            profile_dir,
            old_version,
            new_version,
            output_dir,
            sample_size=sample_size,
            generate_plots=generate_plots,
        )

        if comparison:
            all_comparisons.append(comparison)

    # Create overall summary
    if all_comparisons:
        print("\n" + "=" * 80)
        print("CREATING OVERALL SUMMARY")
        print("=" * 80)

        summary_data = {
            "total_comparisons": len(all_comparisons),
            "version_sequence": versions,
            "comparisons": [
                {
                    "versions": f"{comp['summary']['old_version']} → {comp['summary']['new_version']}",
                    "old_total": comp["summary"]["old_total_records"],
                    "new_total": comp["summary"]["new_total_records"],
                    "record_change": comp["summary"]["record_change"],
                    "added_ids": comp["summary"]["added_ids_count"],
                    "removed_ids": comp["summary"]["removed_ids_count"],
                    "net_change": comp["summary"]["net_change"],
                    "stems_captions_change": comp["stem_changes"]["stems_captions_change"],
                }
                for comp in all_comparisons
            ],
        }

        # Save summary
        summary_path = output_dir / "summary_all_diffs.json"
        with open(summary_path, "w") as f:
            json.dump(summary_data, f, indent=2)
        print(f"Overall summary saved to {summary_path}")

        # Create overall trend plots
        if generate_plots:
            plot_overall_trends(all_comparisons, output_dir)

    print("\n" + "=" * 80)
    print("ALL COMPARISONS COMPLETE")
    print("=" * 80)
    print(f"Results saved to: {output_dir}")


def main():
    """Main entry point for the script."""
    parser = argparse.ArgumentParser(description="Compare Suno data versions and generate diff reports")

    parser.add_argument(
        "--profile-dir",
        type=Path,
        default=Path("/home/vibert/data/suno_data_monitor/outputs"),
        help="Directory containing data profiles",
    )

    parser.add_argument(
        "--output-dir",
        type=Path,
        default=Path("/home/vibert/data/data_diff"),
        help="Directory to save comparison results",
    )

    parser.add_argument(
        "--versions",
        nargs="+",
        default=["v0", "v1", "v2", "v3", "v4", "v5", "v6", "v8", "v9"],
        help="List of versions to compare sequentially (default: v0-v9, excluding v7)",
    )

    parser.add_argument(
        "--sample-size",
        type=int,
        default=20,
        help="Number of samples to extract per category (default: 20)",
    )

    parser.add_argument(
        "--no-plots",
        action="store_true",
        help="Skip generating visualization plots",
    )

    parser.add_argument(
        "--single-pair",
        nargs=2,
        metavar=("OLD_VERSION", "NEW_VERSION"),
        help="Compare only a single version pair (e.g., --single-pair v0 v1)",
    )

    args = parser.parse_args()

    # Validate paths
    if not args.profile_dir.exists():
        print(f"ERROR: Profile directory not found: {args.profile_dir}")
        return 1

    # Single pair comparison
    if args.single_pair:
        old_version, new_version = args.single_pair
        compare_version_pair(
            args.profile_dir,
            old_version,
            new_version,
            args.output_dir,
            sample_size=args.sample_size,
            generate_plots=not args.no_plots,
        )
    else:
        # Run all comparisons
        run_all_diffs(
            args.profile_dir,
            args.output_dir,
            args.versions,
            sample_size=args.sample_size,
            generate_plots=not args.no_plots,
        )

    return 0


if __name__ == "__main__":
    exit(main())