#!/usr/bin/env python3
"""
G5: Assemble Transcript — Ghép các block đã hiệu đính thành file tổng
- Merge theo thứ tự chunk
- Thêm header metadata
- Tạo file final sạch

Usage:
    python3 assemble_transcript.py transcripts/ -o output/012_final.txt
    python3 assemble_transcript.py transcripts/ -o output/012_final.txt --name "012 - Vinaya Page 72-74"
"""

import argparse
import os
from pathlib import Path


def assemble_transcripts(
    input_dir: str,
    output_path: str,
    name: str = None,
    source_audio: str = None,
    separator: str = "\n\n---\n\n"
):
    """
    Assemble all *_clean.txt files into a single final transcript.
    
    Args:
        input_dir: Directory containing *_clean.txt files
        output_path: Output final transcript path
        name: Optional title for the transcript
        source_audio: Optional source audio reference
        separator: Separator between chunks
    """
    txt_files = sorted(Path(input_dir).glob("*_clean.txt"))
    
    if not txt_files:
        print(f"❌ No *_clean.txt files found in {input_dir}")
        return False

    print(f"📦 Assembling {len(txt_files)} chunk(s) from {input_dir}/")
    print()

    parts = []
    total_chars = 0
    
    for txt_path in txt_files:
        with open(txt_path, "r", encoding="utf-8") as f:
            content = f.read().strip()
        
        chunk_name = txt_path.stem.replace("_clean", "")
        chars = len(content)
        total_chars += chars
        print(f"   📄 {txt_path.name}: {chars} chars")
        
        if content:
            parts.append(content)

    # Build final transcript
    final_content = []
    
    # Header
    final_content.append(f"# Transcript: {name or Path(input_dir).parent.name}")
    final_content.append(f"# Chunks: {len(txt_files)} | Total: {total_chars:,} chars")
    if source_audio:
        final_content.append(f"# Source: {source_audio}")
    final_content.append("")
    
    # Body — chunks separated by clear markers
    final_content.append(separator.join(parts))
    final_content.append("")

    os.makedirs(os.path.dirname(output_path) or ".", exist_ok=True)
    
    with open(output_path, "w", encoding="utf-8") as f:
        f.write("\n".join(final_content))

    print(f"\n✅ Final transcript: {output_path}")
    print(f"   Chunks: {len(txt_files)} | Total chars: {total_chars:,}")
    return True


def main():
    parser = argparse.ArgumentParser(
        description="G5: Assemble Transcript — Merge post-processed chunks"
    )
    parser.add_argument("input_dir", help="Directory containing *_clean.txt files")
    parser.add_argument("-o", "--output", required=True,
                        help="Output final transcript path")
    parser.add_argument("--name", default=None,
                        help="Transcript title (e.g. '012 - Vinaya Page 72-74')")
    parser.add_argument("--source", default=None,
                        help="Source audio reference")
    parser.add_argument("--separator", default="\n\n---\n\n",
                        help="Separator between chunks (default: blank + ---)")
    args = parser.parse_args()

    assemble_transcripts(
        args.input_dir,
        args.output,
        args.name,
        args.source,
        args.separator
    )


if __name__ == "__main__":
    main()