#!/usr/bin/env python3
"""
Merge v8 training metadata with vocal captions and pitch range data by ID matching.

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

IMPORTANT: When merging stems_captions, this script properly handles existing captions:
- If a stem key is new, it adds it to the existing stems_captions
- If a stem key already exists, it concatenates the caption lists
- This preserves any existing captions in the base v8 file
"""

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_v8_vocal_data_to_v9_{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_vocal_data_index(vocal_file: Path) -> Dict[str, Dict[str, Any]]:
    """
    Build an index of vocal data by record ID.

    Args:
        vocal_file: Path to vocal captions and pitch JSONL file

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

    vocal_index = {}

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

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

                    if record_id:
                        vocal_index[record_id] = record
                    else:
                        logger.warning(f"Line {line_num}: No ID found in vocal 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 vocal data index: {e}")
        raise

    logger.info(f"Built vocal data index with {len(vocal_index):,} records")
    return vocal_index


def analyze_vocal_data_coverage(vocal_index: Dict[str, Dict[str, Any]]) -> Dict[str, Any]:
    """
    Analyze the coverage of vocal captions and pitch data in the index.

    Args:
        vocal_index: Index of vocal data by ID

    Returns:
        Analysis results
    """
    logger = logging.getLogger(__name__)
    logger.info("Analyzing vocal data coverage...")

    analysis = {
        "total_records": len(vocal_index),
        "records_with_vocal_captions": 0,
        "records_with_pitch_range": 0,
        "total_vocal_captions": 0,
        "vocal_stem_types": set(),
        "captions_per_stem": {},
    }

    for record_id, record in vocal_index.items():
        # Check for vocal captions
        stems_captions = record.get("stems_captions", {})
        has_vocal_captions = False

        for stem_name, captions_list in stems_captions.items():
            if "vocal" in stem_name.lower() or "vox" in stem_name.lower():
                analysis["vocal_stem_types"].add(stem_name)

                for caption_obj in captions_list:
                    if caption_obj.get("prompt_type") == "voice_description_keywords":
                        has_vocal_captions = True
                        analysis["total_vocal_captions"] += 1

        if has_vocal_captions:
            analysis["records_with_vocal_captions"] += 1

        # Check for pitch range
        if "vocal_pitch_range" in record:
            analysis["records_with_pitch_range"] += 1

    # Convert set to sorted list for logging
    analysis["vocal_stem_types"] = sorted(list(analysis["vocal_stem_types"]))

    logger.info(f"Vocal data analysis:")
    logger.info(f"  Total records: {analysis['total_records']:,}")
    logger.info(f"  Records with vocal captions: {analysis['records_with_vocal_captions']:,}")
    logger.info(f"  Records with pitch range: {analysis['records_with_pitch_range']:,}")
    logger.info(f"  Total vocal captions: {analysis['total_vocal_captions']:,}")
    logger.info(f"  Unique vocal stem types: {len(analysis['vocal_stem_types'])}")

    return analysis


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

    Args:
        base_file: Base v8 JSONL file
        vocal_index: Vocal 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,
        "vocal_matches_found": 0,
        "order_violations": [],
        "vocal_ids_in_order": [],
        "errors": [],
    }

    vocal_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 vocal_index:
                        validation_results["vocal_matches_found"] += 1
                        validation_results["vocal_ids_in_order"].append(record_id)

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

                        vocal_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"  Vocal matches found: {validation_results['vocal_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, vocal_index: Dict[str, Dict[str, Any]], output_file: Path, chunk_size: int = 1000
) -> Dict[str, Any]:
    """
    Merge base v8 file with vocal data by ID matching.

    Args:
        base_file: Base v8 JSONL file
        vocal_index: Index of vocal 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,
        "vocal_matches_found": 0,
        "vocal_captions_added": 0,
        "pitch_data_added": 0,
        "records_without_vocal_data": 0,
        "errors": 0,
        "chunks_written": 0,
        "length_issues": 0,  # Track when merged is shorter than original
        "total_chars_input": 0,
        "total_chars_output": 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 vocal data if available
                        if record_id and record_id in vocal_index:
                            stats["vocal_matches_found"] += 1

                            vocal_record = vocal_index[record_id]

                            # Merge vocal stem captions (don't overwrite existing ones)
                            if "stems_captions" in vocal_record:
                                # If base record already has stems_captions, merge them
                                if "stems_captions" in merged_record:
                                    existing_captions = merged_record["stems_captions"]
                                    new_captions = vocal_record["stems_captions"]

                                    # Merge the caption dictionaries
                                    for stem_name, caption_list in new_captions.items():
                                        if stem_name in existing_captions:
                                            # Key exists, concatenate the lists
                                            existing_captions[stem_name].extend(caption_list)
                                        else:
                                            # New key, just add it
                                            existing_captions[stem_name] = caption_list
                                else:
                                    # No existing stems_captions, just use the new ones
                                    merged_record["stems_captions"] = vocal_record["stems_captions"]

                                stats["vocal_captions_added"] += 1

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

                            # Add other vocal-related fields
                            vocal_fields_to_copy = [
                                "pitch_analysis_timestamp",
                                "pitch_analysis_version",
                                "pitch_extraction_method",
                            ]

                            for field in vocal_fields_to_copy:
                                if field in vocal_record:
                                    merged_record[field] = vocal_record[field]
                        else:
                            stats["records_without_vocal_data"] += 1

                        stats["total_output_records"] += 1

                        # Check character lengths
                        original_line_len = len(line.strip())
                        base_record_json = json.dumps(base_record, ensure_ascii=False)
                        base_record_len = len(base_record_json)
                        merged_json = json.dumps(merged_record, ensure_ascii=False)
                        merged_len = len(merged_json)

                        stats["total_chars_input"] += original_line_len
                        stats["total_chars_output"] += merged_len

                        # Log if merged is shorter than original (shouldn't happen)
                        if record_id in vocal_index and merged_len < original_line_len:
                            stats["length_issues"] += 1
                            if stats["length_issues"] <= 10:  # Log first 10 issues
                                logger.warning(
                                    f"Length issue at line {line_num}, ID {record_id}: "
                                    f"original={original_line_len}, base={base_record_len}, "
                                    f"merged={merged_len} (diff={merged_len - original_line_len})"
                                )

                        # Detailed logging for first few records with vocal data
                        if record_id in vocal_index and stats["vocal_matches_found"] <= 5:
                            logger.debug(
                                f"Record {record_id}: original_line={original_line_len}, "
                                f"base_json={base_record_len}, merged={merged_len}, "
                                f"diff={merged_len - original_line_len}"
                            )

                        # Add to chunk buffer
                        chunk_buffer.append(merged_json)

                        # 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], vocal_analysis: Dict[str, Any]
):
    """Print a summary of the merge operation."""
    logger = logging.getLogger(__name__)

    logger.info("\n" + "=" * 80)
    logger.info("MERGE V8 WITH VOCAL DATA SUMMARY")
    logger.info("=" * 80)

    # Validation results
    logger.info("Order Validation:")
    logger.info(f"  Base records checked: {validation_results['base_records_checked']:,}")
    logger.info(f"  Vocal matches found: {validation_results['vocal_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'])}")

    # Vocal data analysis
    logger.info(f"\nVocal Data Analysis:")
    logger.info(f"  Available vocal records: {vocal_analysis['total_records']:,}")
    logger.info(f"  Records with vocal captions: {vocal_analysis['records_with_vocal_captions']:,}")
    logger.info(f"  Records with pitch range: {vocal_analysis['records_with_pitch_range']:,}")
    logger.info(f"  Total vocal captions: {vocal_analysis['total_vocal_captions']:,}")
    logger.info(f"  Unique vocal stem types: {len(vocal_analysis['vocal_stem_types'])}")

    # Merge statistics
    logger.info(f"\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"  Vocal matches found: {stats['vocal_matches_found']:,}")
    logger.info(f"  Records with vocal captions added: {stats['vocal_captions_added']:,}")
    logger.info(f"  Records with pitch data added: {stats['pitch_data_added']:,}")
    logger.info(f"  Records without vocal data: {stats['records_without_vocal_data']:,}")
    logger.info(f"  Errors: {stats['errors']:,}")
    logger.info(f"  Chunks written: {stats['chunks_written']:,}")

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

    # Character length statistics
    if stats["total_chars_input"] > 0:
        char_ratio = (stats["total_chars_output"] / stats["total_chars_input"]) * 100
        logger.info(f"\nCharacter Length Analysis:")
        logger.info(f"  Total input characters: {stats['total_chars_input']:,}")
        logger.info(f"  Total output characters: {stats['total_chars_output']:,}")
        logger.info(f"  Output/Input ratio: {char_ratio:.2f}%")
        logger.info(f"  Length issues detected: {stats['length_issues']:,}")

        if char_ratio < 100:
            logger.warning(f"  WARNING: Output is {100 - char_ratio:.2f}% smaller than input!")

    logger.info("=" * 80)


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

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

    parser.add_argument(
        "--base-file",
        type=str,
        default="/app2/suno/data/auk_v0/metas_v8_tr.jsonl",
        help="Base v8 JSONL file",
    )

    parser.add_argument(
        "--vocal-file",
        type=str,
        default="/home/vibert/data/voice_designer/metas_v6_tr_vocal_captions_pitch_range.jsonl",
        help="JSONL file with vocal captions and pitch range data",
    )

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

    parser.add_argument(
        "--validate-matches",
        type=int,
        default=100,
        help="Number of vocal 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)
    vocal_file = Path(args.vocal_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 vocal_file.exists():
        logger.error(f"Vocal data file not found: {vocal_file}")
        sys.exit(1)

    # Log operation parameters
    logger.info(f"Base file: {base_file}")
    logger.info(f"Vocal data file: {vocal_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 vocal data index
        vocal_index = build_vocal_data_index(vocal_file)

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

        # Analyze vocal data coverage
        vocal_analysis = analyze_vocal_data_coverage(vocal_index)

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

        if not args.skip_validation:
            validation_results = validate_order_alignment(base_file, vocal_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, vocal_index, output_file, args.chunk_size)

        # Print summary
        print_merge_summary(merge_stats, validation_results, vocal_analysis)

        # 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()