#!/usr/bin/env python3 -u
"""
Batch Inference Example for RedenLab ML SDK (Intelligibility & Naturalness)

This script demonstrates how to batch process multiple audio files for
intelligibility or naturalness inference, with automatic chunking for long files.

Usage:
    export REDENLAB_ML_API_KEY='your-api-key'
    python batch_inference.py --input /path/to/audio/folder --output /path/to/output --model-name als-intelligibility-v1

Arguments:
    --input, -i       : Input folder containing audio files
    --output, -o      : Output folder for results
    --model-name, -m  : Model name (e.g., 'als-intelligibility-v1', 'ataxia-naturalness')
    --warmup          : Warm up the backend before processing (recommended for cold starts)
    --workers, -w     : Number of parallel workers for submission (default: 1, sequential)

Available models:
    Intelligibility: ataxia-intelligibility, als-intelligibility-v1, als-intelligibility-v2
    Naturalness:     ataxia-naturalness, als-naturalness-v1, als-naturalness-v2
"""

import argparse
import json
import os
import sys
import threading
from concurrent.futures import ThreadPoolExecutor, as_completed
from contextlib import redirect_stdout
from datetime import datetime
from io import StringIO
from pathlib import Path

import pandas as pd

from redenlab_extract import IntelligibilityClient, NaturalnessClient

# Supported audio extensions
AUDIO_EXTENSIONS = {".wav", ".flac", ".ogg"}

# Valid model names for each client type
INTELLIGIBILITY_MODELS = {
    "ataxia-intelligibility",
    "als-intelligibility-v1",
    "als-intelligibility-v2",
}

NATURALNESS_MODELS = {
    "ataxia-naturalness",
    "als-naturalness-v1",
    "als-naturalness-v2",
}


def get_client_class(model_name: str):
    """Determine the appropriate client class based on model name."""
    if model_name in INTELLIGIBILITY_MODELS:
        return IntelligibilityClient
    elif model_name in NATURALNESS_MODELS:
        return NaturalnessClient
    else:
        raise ValueError(
            f"Unknown model: {model_name}. "
            f"Valid intelligibility models: {INTELLIGIBILITY_MODELS}. "
            f"Valid naturalness models: {NATURALNESS_MODELS}."
        )


def get_score_key(model_name: str) -> str:
    """Get the result score key based on model name."""
    if model_name in INTELLIGIBILITY_MODELS:
        return "intelligibility_score"
    else:
        return "naturalness_score"


def get_audio_files(input_folder: str) -> list[str]:
    """Get list of audio files from input folder."""
    folder = Path(input_folder)
    if not folder.exists():
        raise FileNotFoundError(f"Input folder not found: {input_folder}")

    files = []
    for f in folder.iterdir():
        if f.is_file() and f.suffix.lower() in AUDIO_EXTENSIONS:
            files.append(str(f))

    if not files:
        raise ValueError(
            f"No audio files found in {input_folder}. "
            f"Supported formats: {', '.join(AUDIO_EXTENSIONS)}"
        )

    return sorted(files)


def submit_file(
    client,
    audio_file: str,
) -> tuple[str, list[str], str | None]:
    """
    Submit a file using _submit_with_chunking.

    Returns:
        Tuple of (filename, job_ids, error_message)
    """
    filename = Path(audio_file).name
    error_msg = None

    try:
        job_ids = client._submit_with_chunking(audio_file)
    except Exception as e:
        job_ids = []
        error_msg = str(e)

    return filename, job_ids, error_msg


def main():
    parser = argparse.ArgumentParser(
        description="Batch inference for intelligibility/naturalness using RedenLab ML SDK",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
    python batch_inference.py -i ./audio_files -o ./results -m als-intelligibility-v1
    python batch_inference.py --input /data/recordings --output /data/scores --model-name ataxia-naturalness --warmup
    python batch_inference.py -i ./audio -o ./results -m als-intelligibility-v1 --workers 5
        """,
    )
    parser.add_argument(
        "-i", "--input", required=True, help="Input folder containing audio files"
    )
    parser.add_argument(
        "-o", "--output", required=True, help="Output folder for results"
    )
    parser.add_argument(
        "-m", "--model-name", required=True,
        help="Model name (e.g., 'als-intelligibility-v1', 'ataxia-naturalness')",
    )
    parser.add_argument(
        "--warmup", action="store_true",
        help="Warm up the backend before processing (recommended for cold starts)",
    )
    parser.add_argument(
        "-w", "--workers",
        type=int,
        default=1,
        help="Number of parallel workers for submission (default: 1, sequential)",
    )
    parser.add_argument(
        "--timeout",
        type=int,
        default=3600,
        help="Maximum time to wait for jobs in seconds (default: 3600)",
    )

    args = parser.parse_args()

    # Check for API key
    api_key = os.environ.get("REDENLAB_ML_API_KEY")
    if not api_key:
        print("Error: REDENLAB_ML_API_KEY environment variable not set")
        print("Set it with: export REDENLAB_ML_API_KEY='your-api-key'")
        sys.exit(1)

    # Validate model name and get client class
    try:
        ClientClass = get_client_class(args.model_name)
        score_key = get_score_key(args.model_name)
    except ValueError as e:
        print(f"Error: {e}")
        sys.exit(1)

    # Setup output directory with timestamp and model name
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    run_folder_name = f"{args.model_name}_{timestamp}"

    base_output_path = Path(args.output)
    output_path = base_output_path / run_folder_name
    results_dir = output_path / "json_results"

    output_path.mkdir(parents=True, exist_ok=True)
    results_dir.mkdir(exist_ok=True)

    print("\n" + "=" * 60)
    print("RedenLab ML SDK - Batch Inference")
    print("=" * 60)
    print(f"Input folder:  {args.input}")
    print(f"Output folder: {output_path}")
    print(f"Model:         {args.model_name}")
    print(f"Score key:     {score_key}")
    print(f"Workers:       {args.workers}")
    print(f"Run ID:        {run_folder_name}")
    print("=" * 60)

    # Get audio files
    try:
        audio_files = get_audio_files(args.input)
        print(f"\nFound {len(audio_files)} audio files to process")
    except (FileNotFoundError, ValueError) as e:
        print(f"Error: {e}")
        sys.exit(1)

    # Initialize client
    client = ClientClass(
        api_key=api_key,
        model_name=args.model_name,
        timeout=args.timeout,
    )
    print(f"\nClient initialized: {client}")

    # Warmup if requested
    if args.warmup:
        print("\n" + "-" * 60)
        print("WARMUP: Warming up the backend before processing...")
        print("-" * 60)
        try:
            warmup_result = client.warmup()
            if warmup_result.get("status") == "not_supported":
                print(warmup_result.get("message"))
        except Exception as e:
            print(f"Warning: Warmup failed: {e}")
            print("Continuing with batch processing...")

    # Phase 1: Submit all files using _submit_with_chunking
    # This handles chunking internally and returns list of job_ids per file
    print("\n" + "-" * 60)
    print("PHASE 1: Submitting all files")
    if args.workers > 1:
        print(f"         (using {args.workers} parallel workers)")
        print("         Note: First output appears when initial uploads complete")
    print("-" * 60)

    # Track submissions: filename -> list of job_ids
    file_job_ids: dict[str, list[str]] = {}

    if args.workers == 1:
        # Sequential submission
        for i, audio_file in enumerate(audio_files, 1):
            filename, job_ids, error = submit_file(client, audio_file)

            # Print clean output
            chunk_info = f" ({len(job_ids)} chunks)" if len(job_ids) > 1 else ""
            status = "✗" if error else "✓"
            print(f"  [{i}/{len(audio_files)}] {filename} {status}{chunk_info}")
            if error:
                print(f"         ERROR: {error}")

            file_job_ids[filename] = job_ids
    else:
        # Parallel submission with ThreadPoolExecutor
        print_lock = threading.Lock()
        completed_count = 0

        with ThreadPoolExecutor(max_workers=args.workers) as executor:
            # Submit all tasks
            futures = {
                executor.submit(submit_file, client, audio_file): audio_file
                for audio_file in audio_files
            }

            print(f"\n{args.workers} workers started, processing {len(audio_files)} files...")
            print("(Large files may take several minutes to upload - please wait)\n")
            sys.stdout.flush()

            # Process results as they complete
            for future in as_completed(futures):
                filename, job_ids, error = future.result()

                # Print clean output with completion counter
                with print_lock:
                    completed_count += 1
                    chunk_info = f" ({len(job_ids)} chunks)" if len(job_ids) > 1 else ""
                    status = "✗" if error else "✓"
                    print(f"  [{completed_count}/{len(audio_files)}] {filename} {status}{chunk_info}")
                    if error:
                        print(f"         ERROR: {error}")
                    sys.stdout.flush()

                file_job_ids[filename] = job_ids

    total_jobs = sum(len(ids) for ids in file_job_ids.values())
    print(f"\nSubmitted {total_jobs} total jobs for {len(audio_files)} files")
    sys.stdout.flush()

    # Phase 2: Poll all jobs using poll_multiple for chunked files
    print("\n" + "-" * 60)
    print("PHASE 2: Polling for completion")
    print("-" * 60)
    sys.stdout.flush()

    results: dict[str, dict] = {}

    for i, (filename, job_ids) in enumerate(file_job_ids.items(), 1):
        if not job_ids:
            # Submission failed
            results[filename] = {
                "status": "failed",
                "error": "Submission failed",
                "score": None,
                "std_dev": None,
            }
            print(f"  [{i}/{len(file_job_ids)}] {filename}: SKIPPED (submission failed)")
            continue

        print(f"  [{i}/{len(file_job_ids)}] Polling {filename}...")

        try:
            if len(job_ids) == 1:
                # Single job - use regular poll
                result = client.poll(job_ids[0], timeout=args.timeout)
                score = result.get("result", {}).get(score_key)
                results[filename] = {
                    "status": result.get("status", "unknown"),
                    "score": score,
                    "std_dev": None,
                    "num_chunks": None,
                    "chunk_scores": None,
                    "raw_result": result,
                }
            else:
                # Multiple jobs (chunked) - use poll_multiple which aggregates
                result = client.poll_multiple(job_ids, timeout=args.timeout)
                score = result.get("result", {}).get(score_key)
                std_dev = result.get("result", {}).get("std_dev")
                chunk_scores = result.get("result", {}).get("chunk_scores")
                results[filename] = {
                    "status": result.get("status", "unknown"),
                    "score": score,
                    "std_dev": std_dev,
                    "num_chunks": len(job_ids),
                    "chunk_scores": chunk_scores,
                    "raw_result": result,
                }

            score_str = f"{results[filename]['score']:.4f}" if results[filename]['score'] is not None else "N/A"
            status_icon = "+" if results[filename]["status"] == "completed" else "!"
            print(f"         [{status_icon}] Score: {score_str}")
            sys.stdout.flush()

        except Exception as e:
            print(f"         ERROR: {e}")
            sys.stdout.flush()
            results[filename] = {
                "status": "failed",
                "error": str(e),
                "score": None,
                "std_dev": None,
            }

    # Phase 3: Save results
    print("\n" + "-" * 60)
    print("PHASE 3: Saving results")
    print("-" * 60)
    sys.stdout.flush()

    # Save individual JSON files
    for filename, result in results.items():
        json_filename = Path(filename).stem + ".json"
        json_path = results_dir / json_filename

        # Remove raw_result for cleaner output (it can be verbose)
        result_data = {
            "filename": filename,
            "model_name": args.model_name,
            "score_key": score_key,
            "status": result.get("status"),
            "score": result.get("score"),
            "std_dev": result.get("std_dev"),
            "num_chunks": result.get("num_chunks"),
            "chunk_scores": result.get("chunk_scores"),
        }
        if "error" in result:
            result_data["error"] = result["error"]

        with open(json_path, "w") as f:
            json.dump(result_data, f, indent=2)

    # Create summary DataFrame
    summary_data = []
    for filename, result in results.items():
        summary_data.append({
            "filename": filename,
            "model_name": args.model_name,
            "score": result.get("score"),
            "std_dev": result.get("std_dev"),
            "status": result.get("status"),
            "num_chunks": result.get("num_chunks"),
            "error": result.get("error"),
        })

    summary_df = pd.DataFrame(summary_data)
    summary_path = output_path / "inference_summary.csv"
    summary_df.to_csv(summary_path, index=False)
    print(f"Saved summary to: {summary_path}")

    # Print final summary
    print("\n" + "=" * 60)
    print("SUMMARY")
    print("=" * 60)

    completed = sum(1 for r in results.values() if r.get("status") == "completed")
    partial = sum(1 for r in results.values() if r.get("status") == "partial")
    failed = sum(1 for r in results.values() if r.get("status") == "failed")

    print(f"  Total files:     {len(audio_files)}")
    print(f"  Completed:       {completed}")
    if partial > 0:
        print(f"  Partial:         {partial}")
    print(f"  Failed:          {failed}")
    print()

    # Print scores
    for filename, result in results.items():
        status_icon = "+" if result.get("status") == "completed" else ("!" if result.get("status") == "partial" else "x")
        score_str = f"{result['score']:.4f}" if result.get("score") is not None else "N/A"
        std_str = f" (std: {result['std_dev']:.4f})" if result.get("std_dev") else ""
        chunk_str = f" [{result['num_chunks']} chunks]" if result.get("num_chunks") else ""
        print(f"  [{status_icon}] {filename}: {score_str}{std_str}{chunk_str}")

    print()
    print(f"  Output folder:   {output_path}")
    print(f"  JSON results:    {results_dir}")
    print(f"  Summary CSV:     {summary_path}")
    print("=" * 60)
    sys.stdout.flush()


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print("\n\nInterrupted by user. Exiting...")
        sys.exit(1)
    except Exception as e:
        print(f"\n\nFATAL ERROR: {e}", file=sys.stderr)
        import traceback
        traceback.print_exc()
        sys.exit(1)
