#!/usr/bin/env python3
"""
Merge vocal stem captions with pitch range data by ID matching.

This script processes the base file (w30.jsonl) sequentially and adds pitch data
when a matching ID is found in the pitch file. The pitch file contains a subset
of records from the base file, in order but with gaps.
"""

import argparse
import json
import logging
from datetime import datetime
from pathlib import Path
from typing import Dict, Any, Optional, Set
from tqdm import tqdm
import sys


def setup_logging(output_dir: Path, log_level: str = "INFO") -> logging.Logger:
    """Set up logging configuration."""
    output_dir.mkdir(parents=True, exist_ok=True)
    log_file = output_dir / f"merge_by_id_{datetime.now().strftime('%Y%m%d_%H%M%S')}.log"

    # Create formatter
    formatter = logging.Formatter(
        "%(asctime)s - %(levelname)s - %(message)s", datefmt="%Y-%m-%d %H:%M:%S"
    )

    # File handler
    file_handler = logging.FileHandler(log_file)
    file_handler.setLevel(getattr(logging, log_level.upper()))
    file_handler.setFormatter(formatter)

    # Console handler
    console_handler = logging.StreamHandler(sys.stdout)
    console_handler.setLevel(logging.INFO)
    console_handler.setFormatter(formatter)

    # Create logger
    logger = logging.getLogger(__name__)
    logger.setLevel(getattr(logging, log_level.upper()))
    logger.addHandler(file_handler)
    logger.addHandler(console_handler)

    logger.info(f"Logging to: {log_file}")
    return logger


def build_pitch_index(pitch_file: Path) -> Dict[str, Dict[str, Any]]:
    """
    Build an index of pitch data by record ID.

    Args:
        pitch_file: Path to pitch JSONL file

    Returns:
        Dictionary mapping record ID to pitch data
    """
    logger = logging.getLogger(__name__)
    logger.info(f"Building pitch data index from {pitch_file}")

    pitch_index = {}

    try:
        with open(pitch_file, "r") as f:
            for line_num, line in enumerate(tqdm(f, desc="Indexing pitch data"), 1):
                if not line.strip():
                    continue

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

                    if record_id:
                        pitch_index[record_id] = record
                    else:
                        logger.warning(f"Line {line_num}: No ID found in pitch record")

                except json.JSONDecodeError as e:
                    logger.error(f"Line {line_num}: JSON decode error: {e}")
                    continue

    except Exception as e:
        logger.error(f"Error building pitch index: {e}")
        raise

    logger.info(f"Built pitch index with {len(pitch_index):,} records")
    return pitch_index


def validate_order_alignment(
    base_file: Path, pitch_index: Dict[str, Dict[str, Any]], validation_count: int = 100
) -> Dict[str, Any]:
    """
    Validate that pitch records appear in the same order as base records.

    Args:
        base_file: Base JSONL file
        pitch_index: Pitch data index by ID
        validation_count: Number of matches to validate

    Returns:
        Validation results
    """
    logger = logging.getLogger(__name__)
    logger.info(f"Validating order alignment (checking {validation_count} matches)")

    validation_results = {
        "is_order_preserved": True,
        "base_records_checked": 0,
        "pitch_matches_found": 0,
        "order_violations": [],
        "pitch_ids_in_order": [],
        "errors": [],
    }

    pitch_ids_seen = set()
    matches_validated = 0

    try:
        with open(base_file, "r") as f:
            for line_num, line in enumerate(f, 1):
                if matches_validated >= validation_count:
                    break

                if not line.strip():
                    continue

                try:
                    record = json.loads(line.strip())
                    record_id = record.get("id")
                    validation_results["base_records_checked"] += 1

                    if record_id and record_id in pitch_index:
                        validation_results["pitch_matches_found"] += 1
                        validation_results["pitch_ids_in_order"].append(record_id)

                        # Check if this ID was seen before (would indicate disorder)
                        if record_id in pitch_ids_seen:
                            violation = {
                                "line": line_num,
                                "id": record_id,
                                "error": "Duplicate ID in pitch data",
                            }
                            validation_results["order_violations"].append(violation)
                            validation_results["is_order_preserved"] = False

                        pitch_ids_seen.add(record_id)
                        matches_validated += 1

                except json.JSONDecodeError as e:
                    error_msg = f"Line {line_num}: JSON decode error: {e}"
                    validation_results["errors"].append(error_msg)
                    logger.error(error_msg)
                    continue

    except Exception as e:
        error_msg = f"Error during validation: {e}"
        validation_results["errors"].append(error_msg)
        logger.error(error_msg)
        validation_results["is_order_preserved"] = False

    # Log results
    if validation_results["is_order_preserved"]:
        logger.info(f"✓ Order validation PASSED")
        logger.info(f"  Base records checked: {validation_results['base_records_checked']:,}")
        logger.info(f"  Pitch matches found: {validation_results['pitch_matches_found']:,}")
    else:
        logger.error(
            f"✗ Order validation FAILED with {len(validation_results['order_violations'])} violations"
        )

    return validation_results


def merge_files_by_id(
    base_file: Path, pitch_index: Dict[str, Dict[str, Any]], output_file: Path, chunk_size: int = 1000
) -> Dict[str, Any]:
    """
    Merge base file with pitch data by ID matching.

    Args:
        base_file: Base JSONL file with captions
        pitch_index: Index of pitch data by record ID
        output_file: Output merged file
        chunk_size: Number of records to buffer before writing

    Returns:
        Dictionary with merge statistics
    """
    logger = logging.getLogger(__name__)
    logger.info(f"Starting merge by ID: {base_file} -> {output_file}")

    stats = {
        "total_base_records": 0,
        "total_output_records": 0,
        "pitch_matches_found": 0,
        "pitch_data_added": 0,
        "records_without_pitch": 0,
        "errors": 0,
        "chunks_written": 0,
    }

    # Get total lines for progress bar
    logger.info("Counting total lines in base file...")
    with open(base_file, "r") as f:
        total_lines = sum(1 for _ in f)
    logger.info(f"Total base records: {total_lines:,}")

    try:
        with open(base_file, "r") as f_base, open(output_file, "w") as f_out:
            chunk_buffer = []

            # Progress bar
            with tqdm(total=total_lines, desc="Merging by ID", unit="records") as pbar:
                for line_num, line in enumerate(f_base, 1):
                    if not line.strip():
                        continue

                    try:
                        base_record = json.loads(line.strip())
                        stats["total_base_records"] += 1

                        record_id = base_record.get("id")

                        # Start with base record
                        merged_record = base_record.copy()

                        # Add pitch data if available
                        if record_id and record_id in pitch_index:
                            stats["pitch_matches_found"] += 1

                            pitch_record = pitch_index[record_id]

                            # Add pitch data
                            if "vocal_pitch_range" in pitch_record:
                                merged_record["vocal_pitch_range"] = pitch_record["vocal_pitch_range"]
                                stats["pitch_data_added"] += 1

                            # Add other pitch-related fields
                            pitch_fields_to_copy = [
                                "pitch_analysis_timestamp",
                                "pitch_analysis_version",
                                "pitch_extraction_method",
                            ]

                            for field in pitch_fields_to_copy:
                                if field in pitch_record:
                                    merged_record[field] = pitch_record[field]
                        else:
                            stats["records_without_pitch"] += 1

                        stats["total_output_records"] += 1

                        # Add to chunk buffer
                        chunk_buffer.append(json.dumps(merged_record, ensure_ascii=False))

                        # Write chunk when buffer is full
                        if len(chunk_buffer) >= chunk_size:
                            f_out.write("\n".join(chunk_buffer) + "\n")
                            chunk_buffer = []
                            stats["chunks_written"] += 1

                            # Log progress periodically
                            if stats["chunks_written"] % 100 == 0:
                                logger.info(
                                    f"Processed {stats['total_base_records']:,} records "
                                    f"({stats['chunks_written']} chunks)"
                                )

                        pbar.update(1)

                    except json.JSONDecodeError as e:
                        stats["errors"] += 1
                        logger.error(f"Line {line_num}: JSON decode error: {e}")
                        continue

                    except Exception as e:
                        stats["errors"] += 1
                        logger.error(f"Line {line_num}: Processing error: {e}")
                        continue

                # Write remaining buffer
                if chunk_buffer:
                    f_out.write("\n".join(chunk_buffer) + "\n")
                    stats["chunks_written"] += 1

        logger.info("Merge completed successfully!")

    except Exception as e:
        logger.error(f"Critical error during merge: {e}")
        stats["errors"] += 1
        raise

    return stats


def print_merge_summary(stats: Dict[str, Any], validation_results: Dict[str, Any]):
    """Print a summary of the merge operation."""
    logger = logging.getLogger(__name__)

    logger.info("\n" + "=" * 80)
    logger.info("MERGE BY ID SUMMARY")
    logger.info("=" * 80)

    # Validation results
    logger.info("Order Validation:")
    logger.info(f"  Base records checked: {validation_results['base_records_checked']:,}")
    logger.info(f"  Pitch matches found: {validation_results['pitch_matches_found']:,}")
    logger.info(f"  Order preserved: {'✓ YES' if validation_results['is_order_preserved'] else '✗ NO'}")
    if validation_results["order_violations"]:
        logger.info(f"  Order violations: {len(validation_results['order_violations'])}")

    # Merge statistics
    logger.info("\nMerge Statistics:")
    logger.info(f"  Total base records: {stats['total_base_records']:,}")
    logger.info(f"  Total output records: {stats['total_output_records']:,}")
    logger.info(f"  Pitch matches found: {stats['pitch_matches_found']:,}")
    logger.info(f"  Records with pitch data added: {stats['pitch_data_added']:,}")
    logger.info(f"  Records without pitch: {stats['records_without_pitch']:,}")
    logger.info(f"  Errors: {stats['errors']:,}")
    logger.info(f"  Chunks written: {stats['chunks_written']:,}")

    # Coverage statistics
    if stats["total_base_records"] > 0:
        pitch_coverage = (stats["pitch_matches_found"] / stats["total_base_records"]) * 100
        success_rate = (stats["total_output_records"] / stats["total_base_records"]) * 100
        logger.info(f"\nCoverage:")
        logger.info(f"  Pitch data coverage: {pitch_coverage:.2f}%")
        logger.info(f"  Output success rate: {success_rate:.2f}%")

    logger.info("=" * 80)


def main():
    """Main execution function."""
    parser = argparse.ArgumentParser(
        description="Merge vocal stem captions with pitch range data by ID matching",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
This script processes the base file sequentially and adds pitch data when 
a matching ID is found. The pitch file contains a subset of records from 
the base file, preserving order but with gaps.

Examples:
  # Basic merge with default settings
  python merge_captions_pitch_by_id.py
  
  # Custom validation count and chunk size
  python merge_captions_pitch_by_id.py --validate-matches 200 --chunk-size 2000
        """,
    )

    parser.add_argument(
        "--base-file",
        type=str,
        default="/home/vibert/data/voice_designer/metas_v6_tr_vocal_stems_captioned_w30.jsonl",
        help="Base JSONL file with vocal stem captions",
    )

    parser.add_argument(
        "--pitch-file",
        type=str,
        default="/home/vibert/data/voice_designer/metas_v6_pitch_latest.jsonl",
        help="JSONL file with pitch range data (subset of base)",
    )

    parser.add_argument(
        "--output-file",
        type=str,
        default="/home/vibert/data/voice_designer/metas_v6_tr_vocal_captions_pitch_range.jsonl",
        help="Output merged JSONL file",
    )

    parser.add_argument(
        "--validate-matches",
        type=int,
        default=100,
        help="Number of pitch matches to validate for order preservation",
    )

    parser.add_argument(
        "--chunk-size", type=int, default=1000, help="Number of records to buffer before writing"
    )

    parser.add_argument(
        "--log-dir", type=str, help="Directory for log files (default: same as output file directory)"
    )

    parser.add_argument(
        "--log-level",
        type=str,
        choices=["DEBUG", "INFO", "WARNING", "ERROR"],
        default="INFO",
        help="Logging level",
    )

    parser.add_argument(
        "--skip-validation", action="store_true", help="Skip order validation (faster startup)"
    )

    args = parser.parse_args()

    # Setup paths
    base_file = Path(args.base_file)
    pitch_file = Path(args.pitch_file)
    output_file = Path(args.output_file)

    # Determine log directory
    log_dir = Path(args.log_dir) if args.log_dir else output_file.parent

    # Setup logging
    logger = setup_logging(log_dir, args.log_level)

    # Validate input files exist
    if not base_file.exists():
        logger.error(f"Base file not found: {base_file}")
        sys.exit(1)

    if not pitch_file.exists():
        logger.error(f"Pitch file not found: {pitch_file}")
        sys.exit(1)

    # Log operation parameters
    logger.info(f"Base file: {base_file}")
    logger.info(f"Pitch file: {pitch_file}")
    logger.info(f"Output file: {output_file}")
    logger.info(f"Validation matches: {args.validate_matches}")
    logger.info(f"Chunk size: {args.chunk_size}")

    try:
        # Build pitch data index
        pitch_index = build_pitch_index(pitch_file)

        if not pitch_index:
            logger.error("No pitch data found!")
            sys.exit(1)

        # Validate order unless skipped
        validation_results = {
            "is_order_preserved": True,
            "base_records_checked": 0,
            "pitch_matches_found": 0,
            "order_violations": [],
        }

        if not args.skip_validation:
            validation_results = validate_order_alignment(base_file, pitch_index, args.validate_matches)
        else:
            logger.info("Skipping order validation as requested")

        # Create output directory
        output_file.parent.mkdir(parents=True, exist_ok=True)

        # Perform merge
        merge_stats = merge_files_by_id(base_file, pitch_index, output_file, args.chunk_size)

        # Print summary
        print_merge_summary(merge_stats, validation_results)

        # Verify output file
        if output_file.exists():
            output_size = output_file.stat().st_size
            logger.info(f"Output file created: {output_file} ({output_size:,} bytes)")

            # Quick line count check
            with open(output_file, "r") as f:
                output_lines = sum(1 for _ in f)
            logger.info(f"Output file contains: {output_lines:,} lines")
        else:
            logger.error("Output file was not created!")
            sys.exit(1)

        logger.info("Merge operation completed successfully!")

    except KeyboardInterrupt:
        logger.warning("Operation interrupted by user")
        sys.exit(1)
    except Exception as e:
        logger.error(f"Fatal error: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()