from __future__ import annotations

import argparse
import asyncio
import json
import os
from pathlib import Path
import sys

from dictation.config import (
    DictationConfig,
    MODEL_REPOSITORY,
    MODEL_REVISION,
)


def download_model(config: DictationConfig) -> None:
    from huggingface_hub import snapshot_download

    config.model_dir.mkdir(parents=True, exist_ok=True)
    snapshot_download(
        repo_id=MODEL_REPOSITORY,
        revision=MODEL_REVISION,
        local_dir=config.model_dir,
        allow_patterns=[
            "config.json",
            "model.bin",
            "tokenizer.json",
            "vocabulary.txt",
        ],
    )
    marker = config.model_dir / ".model-revision"
    marker.write_text(f"{MODEL_REPOSITORY}@{MODEL_REVISION}\n", encoding="ascii")
    print(f"Model ready: {config.model_dir}")


def preflight(config: DictationConfig) -> None:
    import ctranslate2

    compute_types = sorted(ctranslate2.get_supported_compute_types("cuda"))
    if config.compute_type not in compute_types:
        raise RuntimeError(
            f"{config.compute_type} is unsupported on CUDA; "
            f"available: {', '.join(compute_types)}"
        )
    result = {
        "cudaDeviceCount": ctranslate2.get_cuda_device_count(),
        "computeTypes": compute_types,
        "selectedComputeType": config.compute_type,
        "modelDir": str(config.model_dir),
        "modelPresent": (config.model_dir / "model.bin").is_file(),
    }
    if result["cudaDeviceCount"] < 1:
        raise RuntimeError("CTranslate2 did not detect a CUDA device")
    print(json.dumps(result, indent=2))


async def transcribe_file(config: DictationConfig, path: Path) -> None:
    from dictation.transcriber import FasterWhisperTranscriber

    transcriber = FasterWhisperTranscriber(
        config.model_dir,
        config.compute_type,
    )
    try:
        from av import open as av_open
        from av.audio.resampler import AudioResampler
        import numpy as np

        samples = []
        resampler = AudioResampler(format="s16", layout="mono", rate=16000)
        with av_open(str(path)) as container:
            for frame in container.decode(audio=0):
                for converted in resampler.resample(frame):
                    samples.append(converted.to_ndarray().reshape(-1))
        if not samples:
            raise RuntimeError("Audio file contains no decodable samples")
        audio = np.concatenate(samples).astype(np.float32) / 32768.0
        result = await transcriber.transcribe(audio, final=True)
        print(
            json.dumps(
                {
                    "text": result.text,
                    "language": result.language,
                    "probability": result.probability,
                },
                ensure_ascii=False,
            )
        )
    finally:
        transcriber.close()


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "command",
        choices=["server", "preflight", "download_model", "transcribe"],
    )
    parser.add_argument("path", nargs="?")
    args = parser.parse_args()
    try:
        config = DictationConfig.from_environment()
        if args.command == "download_model":
            download_model(config)
        elif args.command == "preflight":
            preflight(config)
        elif args.command == "transcribe":
            if not args.path:
                parser.error("transcribe requires an audio file path")
            asyncio.run(transcribe_file(config, Path(args.path)))
        else:
            import uvicorn
            from dictation.server import create_app

            uvicorn.run(
                create_app(config),
                host=config.host,
                port=config.port,
                access_log=False,
            )
    except Exception as error:
        print(f"dictation: {error}", file=sys.stderr)
        raise SystemExit(1) from error


if __name__ == "__main__":
    main()