#!/usr/bin/env python3
"""
Generate self-paired manifests for IndexTTS2 training.
This is a fix for datasets with unique speaker IDs where standard pairing fails.
It creates pairs where Prompt == Target.
"""

import argparse
import json
from pathlib import Path
from typing import List, Dict

def parse_args():
    parser = argparse.ArgumentParser(description="Create self-paired manifests.")
    parser.add_argument("--input", type=Path, required=True, help="Input single-speaker manifest.")
    parser.add_argument("--output", type=Path, required=True, help="Output paired manifest.")
    return parser.parse_args()

def main():
    args = parse_args()
    
    if not args.input.exists():
        raise FileNotFoundError(f"Input manifest not found: {args.input}")
        
    print(f"Processing {args.input} -> {args.output}")
    
    pairs = []
    with args.input.open("r", encoding="utf-8") as f:
        for line in f:
            if not line.strip():
                continue
            entry = json.loads(line)
            
            # Create a self-pair
            # Target fields are from the entry
            # Prompt fields are ALSO from the entry
            
            # Ensure required fields exist
            if "condition_path" not in entry or "emo_vec_path" not in entry:
                print(f"Skipping incomplete entry: {entry.get('id')}")
                continue

            pair = {
                "id": f"{entry['id']}__{entry['id']}", # Unique ID for the pair
                "speaker": entry.get("speaker", "unknown"),
                
                # Prompt (The Voice)
                "prompt_id": entry['id'],
                "prompt_audio_path": entry.get("audio_path", ""),
                "prompt_condition_path": entry["condition_path"],
                "prompt_condition_len": int(entry.get("condition_len", 32)),
                "prompt_emo_vec_path": entry.get("emo_vec_path", ""),
                "prompt_duration": entry.get("duration"),
                "prompt_language": entry.get("language"),
                
                # Target (The Content)
                "target_id": entry['id'],
                "target_audio_path": entry.get("audio_path", ""),
                "target_text": entry.get("text", ""),
                "target_text_ids_path": entry["text_ids_path"],
                "target_text_len": int(entry["text_len"]),
                "target_codes_path": entry["codes_path"],
                "target_code_len": int(entry["code_len"]),
                "target_emo_vec_path": entry.get("emo_vec_path", ""),
                "target_language": entry.get("language"),
            }
            pairs.append(pair)

    # Write output
    args.output.parent.mkdir(parents=True, exist_ok=True)
    with args.output.open("w", encoding="utf-8") as f:
        for p in pairs:
            f.write(json.dumps(p, ensure_ascii=False) + "\n")
            
    print(f"Done. Wrote {len(pairs)} self-paired entries.")

if __name__ == "__main__":
    main()