from __future__ import annotations

import json
import time
from dataclasses import asdict, dataclass
from enum import Enum
from pathlib import Path
from typing import Any, Optional

import typer

from whisper_wrap.data.harvard import HARVARD_SENTENCES

app = typer.Typer(help="Bulk transcription wrapper around faster-whisper.")

DEFAULT_EXTENSIONS = {
    ".mp3",
    ".wav",
    ".m4a",
    ".flac",
    ".ogg",
    ".aac",
    ".wma",
    ".mp4",
    ".mkv",
    ".mov",
    ".avi",
    ".webm",
    ".m4v",
}


class OutputFormat(str, Enum):
    txt = "txt"
    srt = "srt"
    vtt = "vtt"
    json = "json"


class Device(str, Enum):
    cpu = "cpu"
    cuda = "cuda"


class ComputeType(str, Enum):
    float16 = "float16"
    int8 = "int8"


class HarvardFormat(str, Enum):
    text = "text"
    json = "json"


@dataclass
class SegmentRecord:
    start: float
    end: float
    text: str


@dataclass
class FileRecord:
    file: str
    duration: Optional[float]
    transcription_time: Optional[float]
    model: str
    language: Optional[str]
    device: str
    compute_type: str
    segments: list[SegmentRecord]


def normalize_extensions(ext_values: Optional[list[str]]) -> set[str]:
    if not ext_values:
        return DEFAULT_EXTENSIONS
    normalized: set[str] = set()
    for ext in ext_values:
        if not ext:
            continue
        ext = ext.strip().lower()
        if not ext.startswith("."):
            ext = f".{ext}"
        normalized.add(ext)
    return normalized or DEFAULT_EXTENSIONS


def harvard_attribution() -> list[str]:
    return [
        "Harvard Sentences (1965 revised list).",
        "Source: https://www.cs.columbia.edu/~hgs/audio/harvard.html",
        "From \"IEEE Recommended Practice for Speech Quality Measurements.\"",
        "IEEE Transactions on Audio and Electroacoustics, Vol. 17, Issue 3, 1969.",
    ]


def discover_files(input_dir: Path, recursive: bool, extensions: set[str]) -> list[Path]:
    glob_pattern = "**/*" if recursive else "*"
    files = [
        p
        for p in input_dir.glob(glob_pattern)
        if p.is_file() and p.suffix.lower() in extensions
    ]
    return sorted(files)


def format_seconds(seconds: float, vtt: bool = False) -> str:
    if seconds < 0:
        seconds = 0
    total_ms = int(round(seconds * 1000))
    hours = total_ms // 3_600_000
    rem = total_ms % 3_600_000
    minutes = rem // 60_000
    rem %= 60_000
    secs = rem // 1000
    ms = rem % 1000
    sep = "." if vtt else ","
    return f"{hours:02d}:{minutes:02d}:{secs:02d}{sep}{ms:03d}"


def metadata_lines(record: FileRecord) -> list[str]:
    duration_str = (
        f"{record.duration:.2f}s" if isinstance(record.duration, (int, float)) else "unknown"
    )
    transcribe_str = (
        f"{record.transcription_time:.2f}s"
        if isinstance(record.transcription_time, (int, float))
        else "unknown"
    )
    return [
        f"file: {record.file}",
        f"duration: {duration_str}",
        f"transcription_time: {transcribe_str}",
        f"model: {record.model}",
        f"language: {record.language or 'auto'}",
        f"device: {record.device}",
        f"compute_type: {record.compute_type}",
    ]


def write_txt(path: Path, record: FileRecord, include_metadata: bool = True) -> None:
    body = "\n".join(seg.text.strip() for seg in record.segments if seg.text.strip())
    if include_metadata:
        header = "\n".join([f"# {line}" for line in metadata_lines(record)])
        path.write_text(f"{header}\n\n{body}\n", encoding="utf-8")
    else:
        path.write_text(f"{body}\n", encoding="utf-8")


def write_srt(path: Path, record: FileRecord, start_idx: int = 1, offset: float = 0.0) -> int:
    lines: list[str] = []
    # SRT doesn't support NOTE headers; keep metadata out of SRT.

    idx = start_idx
    for seg in record.segments:
        start = format_seconds(seg.start + offset, vtt=False)
        end = format_seconds(seg.end + offset, vtt=False)
        text = seg.text.strip()
        if not text:
            continue
        lines.extend([str(idx), f"{start} --> {end}", text, ""])
        idx += 1

    mode = "a" if path.exists() and start_idx != 1 else "w"
    with path.open(mode, encoding="utf-8") as f:
        if mode == "a":
            f.write("\n")
        f.write("\n".join(lines).rstrip())
        f.write("\n")
    return idx


def write_vtt(
    path: Path,
    record: FileRecord,
    offset: float = 0.0,
    append: bool = False,
    include_metadata: bool = True,
) -> None:
    lines: list[str] = []
    if not append:
        lines.append("WEBVTT")
        lines.append("")
        if include_metadata:
            lines.extend([f"NOTE {line}" for line in metadata_lines(record)])
            lines.append("")

    for seg in record.segments:
        start = format_seconds(seg.start + offset, vtt=True)
        end = format_seconds(seg.end + offset, vtt=True)
        text = seg.text.strip()
        if not text:
            continue
        lines.extend([f"{start} --> {end}", text, ""])

    mode = "a" if append and path.exists() else "w"
    with path.open(mode, encoding="utf-8") as f:
        if mode == "a":
            f.write("\n")
        f.write("\n".join(lines).rstrip())
        f.write("\n")


def write_json(path: Path, record: FileRecord, include_metadata: bool = True) -> None:
    payload: dict[str, Any] = {
        "segments": [asdict(seg) for seg in record.segments],
    }
    if include_metadata:
        payload["metadata"] = {
            "file": record.file,
            "duration": record.duration,
            "transcription_time": record.transcription_time,
            "model": record.model,
            "language": record.language,
            "device": record.device,
            "compute_type": record.compute_type,
        }
    path.write_text(json.dumps(payload, indent=2, ensure_ascii=False) + "\n", encoding="utf-8")


def write_file_output(path: Path, fmt: OutputFormat, record: FileRecord, include_metadata: bool) -> None:
    if fmt == OutputFormat.txt:
        write_txt(path, record, include_metadata=include_metadata)
    elif fmt == OutputFormat.srt:
        write_srt(path, record)
    elif fmt == OutputFormat.vtt:
        write_vtt(path, record, include_metadata=include_metadata)
    elif fmt == OutputFormat.json:
        write_json(path, record, include_metadata=include_metadata)
    else:
        raise ValueError(f"Unsupported output format: {fmt}")


def write_merged_output(
    path: Path,
    fmt: OutputFormat,
    records: list[FileRecord],
    include_metadata: bool,
) -> None:
    if not records:
        return

    if fmt == OutputFormat.txt:
        chunks: list[str] = []
        for rec in records:
            body = "\n".join(seg.text.strip() for seg in rec.segments if seg.text.strip())
            if include_metadata:
                header = "\n".join([f"# {line}" for line in metadata_lines(rec)])
                chunks.append(f"{header}\n\n{body}".strip())
            else:
                chunks.append(body.strip())
        path.write_text("\n\n".join(chunks) + "\n", encoding="utf-8")
        return

    if fmt == OutputFormat.json:
        payload: dict[str, Any] = {
            "files": [
                {
                    **(
                        {
                            "metadata": {
                                "file": rec.file,
                                "duration": rec.duration,
                                "transcription_time": rec.transcription_time,
                                "model": rec.model,
                                "language": rec.language,
                                "device": rec.device,
                                "compute_type": rec.compute_type,
                            }
                        }
                        if include_metadata
                        else {}
                    ),
                    "segments": [asdict(seg) for seg in rec.segments],
                }
                for rec in records
            ],
        }
        if include_metadata:
            payload["metadata"] = {
                "merged": True,
                "files": len(records),
            }
        path.write_text(json.dumps(payload, indent=2, ensure_ascii=False) + "\n", encoding="utf-8")
        return

    if fmt == OutputFormat.srt:
        if path.exists():
            path.unlink()
        idx = 1
        offset = 0.0
        for rec in records:
            idx = write_srt(path, rec, start_idx=idx, offset=offset)
            if rec.segments:
                offset += rec.segments[-1].end
        return

    if fmt == OutputFormat.vtt:
        if path.exists():
            path.unlink()
        offset = 0.0
        for i, rec in enumerate(records):
            write_vtt(path, rec, offset=offset, append=i > 0, include_metadata=include_metadata)
            if rec.segments:
                offset += rec.segments[-1].end
        return

    raise ValueError(f"Unsupported output format: {fmt}")


@app.command()
def harvard(
    format: HarvardFormat = typer.Option(
        HarvardFormat.text, "--format", help="Output format: text or json."
    ),
    limit: Optional[int] = typer.Option(
        None, "--limit", min=1, help="Limit the number of sentences returned."
    ),
    attribution: bool = typer.Option(
        True,
        "--attribution/--no-attribution",
        help="Include source attribution in output.",
    ),
) -> None:
    sentences = list(HARVARD_SENTENCES)
    if limit:
        sentences = sentences[:limit]

    if format == HarvardFormat.json:
        payload: dict[str, Any] = {
            "sentences": sentences,
        }
        if attribution:
            payload["attribution"] = harvard_attribution()
        typer.echo(json.dumps(payload, indent=2, ensure_ascii=False))
        return

    lines: list[str] = []
    if attribution:
        lines.extend([f"# {line}" for line in harvard_attribution()])
        lines.append("")
    lines.extend(sentences)
    typer.echo("\n".join(lines))


@app.command()
def main(
    input_dir: Path = typer.Argument(..., exists=True, file_okay=False, dir_okay=True),
    model: str = typer.Option("small", "--model", help="Whisper model size or path."),
    language: Optional[str] = typer.Option(
        None, "--language", help="Language code (optional, auto-detect if omitted)."
    ),
    device: Device = typer.Option(Device.cpu, "--device", help="Inference device."),
    compute_type: ComputeType = typer.Option(
        ComputeType.float16, "--compute-type", help="Computation precision/type."
    ),
    threads: Optional[int] = typer.Option(
        None, "--threads", min=1, help="CPU threads for inference."
    ),
    output_dir: Path = typer.Option(
        Path("outputs"), "--output-dir", help="Directory for transcription outputs."
    ),
    merge: bool = typer.Option(False, "--merge", help="Also produce one merged output file."),
    fmt: OutputFormat = typer.Option(OutputFormat.txt, "--format", help="Output format."),
    metadata: bool = typer.Option(True, "--metadata/--no-metadata", help="Include metadata headers."),
    ext: Optional[list[str]] = typer.Option(
        None, "--ext", help="Extension filter, repeatable (e.g., --ext mp3 --ext wav)."
    ),
    recursive: bool = typer.Option(False, "--recursive", help="Recurse into subdirectories."),
) -> None:
    from faster_whisper import WhisperModel

    extensions = normalize_extensions(ext)
    files = discover_files(input_dir, recursive=recursive, extensions=extensions)

    if not files:
        typer.secho("No matching media files found.", fg=typer.colors.YELLOW)
        raise typer.Exit(code=1)

    output_dir.mkdir(parents=True, exist_ok=True)

    if device == Device.cpu and compute_type == ComputeType.float16:
        typer.secho(
            "CPU does not support float16 efficiently. Switching compute type to int8.",
            fg=typer.colors.YELLOW,
        )
        compute_type = ComputeType.int8

    typer.secho(
        f"Loading model '{model}' on {device.value} ({compute_type.value})...",
        fg=typer.colors.CYAN,
    )
    model_kwargs: dict[str, Any] = {
        "device": device.value,
        "compute_type": compute_type.value,
    }
    if threads is not None:
        model_kwargs["cpu_threads"] = threads

    whisper_model = WhisperModel(model, **model_kwargs)

    records: list[FileRecord] = []
    failures: list[tuple[Path, str]] = []

    total = len(files)
    overall_start = time.perf_counter()
    for idx, media_path in enumerate(files, start=1):
        typer.echo(f"[{idx}/{total}] Transcribing: {media_path.name}")
        file_start = time.perf_counter()
        try:
            segments_iter, info = whisper_model.transcribe(
                str(media_path),
                language=language,
            )
            segments = [
                SegmentRecord(start=s.start, end=s.end, text=s.text)
                for s in segments_iter
            ]
            detected_lang = language or getattr(info, "language", None)
            duration = getattr(info, "duration", None)

            elapsed = time.perf_counter() - file_start
            record = FileRecord(
                file=str(media_path),
                duration=duration,
                transcription_time=elapsed,
                model=model,
                language=detected_lang,
                device=device.value,
                compute_type=compute_type.value,
                segments=segments,
            )
            records.append(record)

            out_path = output_dir / f"{media_path.stem}.{fmt.value}"
            write_file_output(out_path, fmt, record, include_metadata=metadata)
            typer.echo(f"  ✓ Done in {elapsed:.2f}s")
        except Exception as exc:  # noqa: BLE001
            failures.append((media_path, str(exc)))
            typer.secho(f"  Error: {exc}", fg=typer.colors.RED)

    total_elapsed = time.perf_counter() - overall_start

    if merge and records:
        merged_path = output_dir / f"merged.{fmt.value}"
        write_merged_output(merged_path, fmt, records, include_metadata=metadata)

    typer.echo("\nSummary")
    typer.echo(f"- Total files scanned: {total}")
    typer.echo(f"- Successful transcriptions: {len(records)}")
    typer.echo(f"- Failed transcriptions: {len(failures)}")
    typer.echo(f"- Total time: {total_elapsed:.2f}s")
    typer.echo(f"- Output directory: {output_dir.resolve()}")
    if merge and records:
        typer.echo(f"- Merged output: {(output_dir / f'merged.{fmt.value}').resolve()}")

    if failures:
        typer.secho("\nErrors", fg=typer.colors.RED)
        for failed_path, err in failures:
            typer.echo(f"- {failed_path}: {err}")

    if records and not failures:
        raise typer.Exit(code=0)
    if records and failures:
        raise typer.Exit(code=2)
    raise typer.Exit(code=1)


if __name__ == "__main__":
    app()
