#!/usr/bin/env python3
"""
Pipeline Orchestrator — Điều phối toàn bộ quy trình 5 giai đoạn
- Phát hiện file audio mới → tự động chạy pipeline
- Xử lý cuốn chiếu: chunk → transcribe → post-process → assemble
- State tracking: ghi log từng bước, resume nếu fail
- Thiết kế để chạy qua cron job

Usage:
    # Full pipeline
    python3 pipeline_orchestrator.py 012.mp3 --name 012

    # Resume từ giai đoạn cụ thể
    python3 pipeline_orchestrator.py 012.mp3 --name 012 --resume-from g4

    # Chỉ chạy 1 giai đoạn
    python3 pipeline_orchestrator.py 012.mp3 --name 012 --only g3

    # Dry run (kiểm tra cấu hình)
    python3 pipeline_orchestrator.py 012.mp3 --name 012 --dry-run
"""

import argparse
import json
import os
import subprocess
import sys
import time
from datetime import datetime
from pathlib import Path


# ─── Configuration ───
SCRIPTS_DIR = Path(__file__).parent
OUTPUT_DIR = Path(__file__).parent.parent / "output"

STAGES = {
    "g1": {
        "name": "Audio Preprocessing",
        "script": "audio_preprocess.py",
        "description": "Chuẩn hóa → 16kHz Mono + Compand filter",
    },
    "g2": {
        "name": "Audio Chunking",
        "script": "audio_chunk.py",
        "description": "Cắt thành block 2 phút",
    },
    "g3": {
        "name": "Gemini Transcribe",
        "script": "gemini_transcribe.py",
        "description": "Gemini 3 Flash Preview → TXT + SRT (1 API call, có timestamp)",
    },
    "g5": {
        "name": "Assemble Transcript",
        "script": "assemble_transcript.py",
        "description": "Ghép các block thành file tổng",
    },
}


class PipelineState:
    """Track pipeline progress with state file"""

    def __init__(self, state_path: str):
        self.state_path = state_path
        self.data = self._load()

    def _load(self) -> dict:
        if os.path.exists(self.state_path):
            with open(self.state_path, "r") as f:
                return json.load(f)
        return {
            "pipeline_version": "1.0",
            "stages": {},
            "runs": []
        }

    def save(self):
        os.makedirs(os.path.dirname(self.state_path), exist_ok=True)
        with open(self.state_path, "w") as f:
            json.dump(self.data, f, ensure_ascii=False, indent=2)

    def stage_start(self, stage: str):
        self.data["stages"][stage] = {
            "status": "running",
            "started_at": datetime.now().isoformat(),
            "completed_at": None,
            "error": None
        }
        self.save()

    def stage_complete(self, stage: str):
        self.data["stages"][stage] = {
            "status": "completed",
            "started_at": self.data["stages"].get(stage, {}).get("started_at"),
            "completed_at": datetime.now().isoformat(),
            "error": None
        }
        self.save()

    def stage_fail(self, stage: str, error: str):
        self.data["stages"][stage] = {
            "status": "failed",
            "started_at": self.data["stages"].get(stage, {}).get("started_at"),
            "completed_at": None,
            "error": str(error)[:500]
        }
        self.save()

    def is_completed(self, stage: str) -> bool:
        return self.data.get("stages", {}).get(stage, {}).get("status") == "completed"

    def is_failed(self, stage: str) -> bool:
        return self.data.get("stages", {}).get(stage, {}).get("status") == "failed"

    def run_start(self, audio_file: str):
        self.data["runs"].append({
            "audio_file": audio_file,
            "started_at": datetime.now().isoformat(),
            "completed_at": None,
        })
        self.data["stages"] = {}  # Reset stages for new run
        self.save()

    def run_complete(self):
        if self.data["runs"]:
            self.data["runs"][-1]["completed_at"] = datetime.now().isoformat()


def run_script(script_name: str, args: list, stage_id: str) -> bool:
    """
    Run a pipeline script as subprocess.
    
    Returns:
        True if script exits with code 0
    """
    script_path = SCRIPTS_DIR / script_name
    cmd = [sys.executable, str(script_path)] + args

    print(f"   🏃 Running: {' '.join(cmd)}")
    print()

    try:
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=3600)
        print(result.stdout)
        if result.returncode != 0:
            print(f"   ❌ STDERR:\n{result.stderr[-1000:]}")
            return False
        return True
    except subprocess.TimeoutExpired:
        print(f"   ❌ Timeout (>1h) for {stage_id}")
        return False
    except Exception as e:
        print(f"   ❌ Exception: {e}")
        return False


def run_pipeline(
    audio_path: str,
    name: str,
    chunk_duration: int = 300,
    start_from: str = "g1",
    only_stage: str = None,
    dry_run: bool = False
):
    """
    Execute the full pipeline.
    
    Args:
        audio_path: Path to input audio
        name: Project name (e.g. '012')
        chunk_duration: Chunk duration in seconds (default: 300 = 5 min)
        start_from: Stage to start/resume from
        only_stage: Run only this stage and stop
        dry_run: Print what would be done without executing
    """
    print("=" * 60)
    print("🏗️  Myanmar Audio → Text Pipeline")
    print("=" * 60)
    print(f"   Audio: {audio_path}")
    print(f"   Name: {name}")
    print(f"   Chunk: {chunk_duration}s ({chunk_duration//60} min)")
    if dry_run:
        print(f"   Mode: DRY RUN")
    print()

    # Paths
    clean_audio = OUTPUT_DIR / f"{name}_clean.wav"
    chunks_dir = OUTPUT_DIR / "chunks"
    transcripts_dir = OUTPUT_DIR / "transcripts"
    final_output = OUTPUT_DIR / f"{name}_final.txt"
    state_file = OUTPUT_DIR / f"{name}_state.json"

    state = PipelineState(str(state_file))
    
    if not dry_run:
        state.run_start(audio_path)

    # Determine which stages to run
    stage_ids = list(STAGES.keys())
    if only_stage:
        if only_stage not in STAGES:
            print(f"❌ Unknown stage: {only_stage}. Must be one of: {', '.join(STAGES)}")
            sys.exit(1)
        stage_ids = [only_stage]
    else:
        start_idx = stage_ids.index(start_from) if start_from in STAGES else 0
        stage_ids = stage_ids[start_idx:]

    success = True

    for stage_id in stage_ids:
        stage = STAGES[stage_id]
        print(f"\n{'─' * 50}")
        print(f"📌 {stage_id.upper()}: {stage['name']}")
        print(f"   {stage['description']}")
        print(f"{'─' * 50}")

        if not dry_run and state.is_completed(stage_id) and not only_stage:
            print(f"   ⏭️  Already completed, skipping")
            continue

        if not dry_run:
            state.stage_start(stage_id)

        # Build stage-specific arguments
        stage_args = []
        ok = True

        if stage_id == "g1":
            stage_args = [
                audio_path,
                "-o", str(clean_audio),
            ]

        elif stage_id == "g2":
            stage_args = [
                str(clean_audio),
                "-d", str(chunk_duration),
                "-o", str(chunks_dir),
            ]

        elif stage_id == "g3":
            if not any((chunks_dir).glob("*.wav")):
                print(f"   ⚠️  No chunks found in {chunks_dir}, run G2 first")
                ok = False
            else:
                stage_args = [
                    str(chunks_dir),
                    "-o", str(transcripts_dir),
                ]

        elif stage_id == "g5":
            if not any((transcripts_dir).glob("*_gemini.txt")):
                print(f"   ⚠️  No gemini transcripts found, run G3 first")
                ok = False
            else:
                stage_args = [
                    str(transcripts_dir),
                    "-o", str(final_output),
                    "--name", name,
                    "--source", os.path.basename(audio_path),
                ]

        if not ok:
            if not dry_run:
                state.stage_fail(stage_id, "Missing input files from previous stage")
            success = False
            if only_stage:
                break
            continue

        if dry_run:
            print(f"   🏁 Would run: {stage['script']} {' '.join(stage_args)}")
        else:
            ok = run_script(stage["script"], stage_args, stage_id)

        if ok:
            if not dry_run:
                state.stage_complete(stage_id)
        else:
            if not dry_run:
                state.stage_fail(stage_id, "Script returned non-zero exit code")
            success = False
            if only_stage:
                break
            # Continue to next stage anyway (partial output may still be useful)

    # Summary
    print(f"\n{'=' * 60}")
    if success:
        print(f"✅ Pipeline completed successfully!")
        if not dry_run:
            state.run_complete()
            print(f"   Final output: {final_output}")
    else:
        print(f"⚠️  Pipeline completed with errors. Check state file:")
        print(f"   {state_file}")
    print(f"{'=' * 60}")

    return success


def main():
    parser = argparse.ArgumentParser(
        description="🏗️  Myanmar Audio → Text Pipeline Orchestrator"
    )
    parser.add_argument("audio", help="Input audio file (.mp3, .wav, .m4a...)")
    parser.add_argument("--name", required=True,
                        help="Project name (e.g. '012')")
    parser.add_argument("--chunk-duration", type=int, default=300,
                        help="Chunk duration in seconds (default: 300 = 5 min)")
    parser.add_argument("--resume-from", default="g1",
                        choices=list(STAGES.keys()),
                        help="Resume pipeline from this stage")
    parser.add_argument("--only", default=None,
                        choices=list(STAGES.keys()),
                        help="Run only this specific stage")
    parser.add_argument("--dry-run", action="store_true",
                        help="Show what would be done without executing")
    args = parser.parse_args()

    if not os.path.exists(args.audio):
        print(f"❌ Audio file not found: {args.audio}")
        sys.exit(1)

    run_pipeline(
        args.audio,
        args.name,
        args.chunk_duration,
        args.resume_from,
        args.only,
        args.dry_run
    )


if __name__ == "__main__":
    main()