from __future__ import annotations

import argparse
import sys
from pathlib import Path

from transcriber_core import DEFAULT_CHUNK_SECONDS, DEFAULT_MODEL_DIR
from transcriber_core import GemmaTranscriber, TranscriberSettings


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Transcribe an audio or video file with a local Gemma model."
    )
    parser.add_argument(
        "audio_file",
        type=Path,
        help="Path to the audio or video file.",
    )
    parser.add_argument(
        "--model-dir",
        type=Path,
        default=DEFAULT_MODEL_DIR,
        help=f"Path to the local Gemma model. Default: {DEFAULT_MODEL_DIR}",
    )
    parser.add_argument(
        "--language",
        default="ja",
        help="Target transcript language, for example ja or en. Default: ja",
    )
    parser.add_argument(
        "--mode",
        choices=["speech", "song"],
        default="speech",
        help="Transcription mode. Use song for singing vocals. Default: speech",
    )
    parser.add_argument(
        "--chunk-seconds",
        type=int,
        default=DEFAULT_CHUNK_SECONDS,
        help="Audio chunk length in seconds. Default: 30",
    )
    parser.add_argument(
        "--max-new-tokens",
        type=int,
        default=256,
        help="Maximum tokens to generate for each chunk. Default: 256",
    )
    parser.add_argument(
        "--dtype",
        choices=["auto", "bfloat16", "float16", "float32"],
        default="auto",
        help="Torch dtype for model loading. Default: auto",
    )
    parser.add_argument(
        "--device-map",
        default="auto",
        help="Transformers device_map value. Default: auto",
    )
    parser.add_argument(
        "--max-memory",
        default='{"cuda:0":"10GiB","cuda:1":"10GiB","cpu":"32GiB"}',
        help='Optional max_memory JSON for multi-GPU loading. '
        'Example: {"cuda:0":"10GiB","cuda:1":"10GiB","cpu":"32GiB"}',
    )
    parser.add_argument(
        "--ffmpeg-path",
        default="",
        help="Optional full path to ffmpeg.exe for video transcription.",
    )
    parser.add_argument(
        "--output",
        type=Path,
        help="Optional output text file path.",
    )
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    settings = TranscriberSettings(
        model_dir=args.model_dir,
        language=args.language,
        mode=args.mode,
        chunk_seconds=args.chunk_seconds,
        max_new_tokens=args.max_new_tokens,
        dtype=args.dtype,
        device_map=args.device_map,
        max_memory=args.max_memory,
        ffmpeg_path=args.ffmpeg_path,
    )

    try:
        print("Loading model and processor...", file=sys.stderr)
        transcriber = GemmaTranscriber(settings)
        final_text = transcriber.transcribe(
            args.audio_file,
            progress_callback=lambda message: print(message, file=sys.stderr),
        )
    except RuntimeError as exc:
        raise SystemExit(str(exc)) from exc

    print(final_text)

    if args.output:
        args.output.write_text(final_text, encoding="utf-8")
        print(f"Saved transcript to: {args.output}", file=sys.stderr)

    return 0


if __name__ == "__main__":
    raise SystemExit(main())
