"""Orchestrate convert -> separate -> transcribe.

Two modes:
  - sequential: for each file, run all enabled stages before moving to the next.
  - staged:     run stage 1 over all files, then stage 2 over all files, then stage 3.

Intermediate artifacts (convert.wav, vocals.wav) are written to a temp dir and
NEVER persisted to the output directory — unless `keep_intermediate=True`.
Only the artifact of the LAST enabled stage is moved into `<output>/<stem>/`.
"""
from __future__ import annotations

import shutil
import tempfile
import threading
from dataclasses import dataclass
from pathlib import Path
from typing import Callable, Optional

from .config import PipelineConfig
from .convert import convert_to_wav
from .separate import separate_vocals
from .transcribe import transcribe


AUDIO_EXTENSIONS = {
    ".mp3", ".wav", ".flac", ".ogg", ".m4a", ".aac", ".opus", ".wma",
}

LogFn = Callable[[str], None]
ProgressFn = Callable[[int, int, str], None]   # done, total, stage


@dataclass
class FileResult:
    input_path: Path
    text: Optional[str]
    output_dir: Optional[Path]
    error: Optional[str] = None


class Cancelled(Exception):
    """Raised when the stop event is set."""


# ---------------------------------------------------------------- helpers

_ARTIFACT_NAMES = {
    "convert": "convert.wav",
    "separate": "vocals.wav",
    "transcribe": "text.txt",
}


def collect_audio_files(input_dir: Path, recursive: bool) -> list[Path]:
    it = input_dir.rglob("*") if recursive else input_dir.iterdir()
    return sorted(p for p in it if p.is_file() and p.suffix.lower() in AUDIO_EXTENSIONS)


def _final_stage(config: PipelineConfig) -> Optional[str]:
    """The stage whose artifact is the final result (single element)."""
    if config.transcribe.enabled:
        return "transcribe"
    if config.separate.enabled:
        return "separate"
    if config.convert.enabled:
        return "convert"
    return None


def _check_cancel(stop_event: Optional[threading.Event]) -> None:
    if stop_event is not None and stop_event.is_set():
        raise Cancelled()


def _persist_final(work_dir: Path, persist_dir: Path, final_stage: str, log: LogFn) -> None:
    """Move the final-stage artifact from work_dir into persist_dir."""
    name = _ARTIFACT_NAMES[final_stage]
    src = work_dir / name
    if not src.exists():
        raise FileNotFoundError(f"Final artifact missing: {src}")
    persist_dir.mkdir(parents=True, exist_ok=True)
    dst = persist_dir / name
    if dst.exists():
        dst.unlink()
    shutil.move(str(src), str(dst))
    log(f"    [persist] {final_stage} -> {persist_dir.name}/{name}")


def _work_dir_for(
    config: PipelineConfig,
    stem: str,
    output_root: Path,
    tmp_root: Optional[Path],
) -> tuple[Path, Path]:
    """Return (work_dir, persist_dir).

    - keep_intermediate=True  -> work_dir == persist_dir == output_root/stem
    - keep_intermediate=False -> work_dir = tmp_root/stem, persist_dir = output_root/stem
    """
    persist_dir = output_root / stem
    if config.keep_intermediate:
        persist_dir.mkdir(parents=True, exist_ok=True)
        return persist_dir, persist_dir
    if tmp_root is None:
        raise ValueError("tmp_root required when keep_intermediate=False")
    work_dir = tmp_root / stem
    work_dir.mkdir(parents=True, exist_ok=True)
    return work_dir, persist_dir


# ---------------------------------------------------------------- entry

def run_pipeline(
    config: PipelineConfig,
    log: LogFn = print,
    on_progress: Optional[ProgressFn] = None,
    stop_event: Optional[threading.Event] = None,
) -> list[FileResult]:
    input_dir = Path(config.input_dir)
    if not input_dir.is_dir():
        raise FileNotFoundError(f"Input directory not found: {input_dir}")

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

    files = collect_audio_files(input_dir, config.recursive)
    if not files:
        log("Аудиофайлы не найдены.")
        return []

    mode = (config.mode or "sequential").lower()
    log(f"Найдено треков: {len(files)}   |   режим: {mode}   |   "
        f"keep_intermediate={config.keep_intermediate}")
    if not config.keep_intermediate:
        log("  Промежуточные файлы будут во временной папке и удалены после прогона.")

    if mode == "staged":
        return _run_staged(config, files, log, on_progress, stop_event)
    return _run_sequential(config, files, log, on_progress, stop_event)


# ---------------------------------------------------------------- sequential

def _run_sequential(
    config: PipelineConfig,
    files: list[Path],
    log: LogFn,
    on_progress: Optional[ProgressFn],
    stop_event: Optional[threading.Event],
) -> list[FileResult]:
    total = len(files)
    final_stage = _final_stage(config)
    results: list[FileResult] = []
    output_root = Path(config.output_dir)

    for idx, f in enumerate(files, 1):
        if stop_event is not None and stop_event.is_set():
            log("⏹ Отменено")
            break
        if on_progress:
            on_progress(idx - 1, total, "file")
        log(f"\n▶ [{idx}/{total}] {f.name}")

        tmp_root = Path(tempfile.mkdtemp(prefix=f"at_{f.stem}_"))
        work_dir, persist_dir = _work_dir_for(config, f.stem, output_root, tmp_root)

        audio_path: Path = f
        text: Optional[str] = None
        error: Optional[str] = None
        processed_ok = False

        try:
            if config.convert.enabled:
                _check_cancel(stop_event)
                log("  [1] Конвертация в WAV...")
                target = work_dir / "convert.wav"
                if config.overwrite or not target.exists():
                    convert_to_wav(
                        f, target,
                        sample_rate=config.convert.sample_rate,
                        channels=config.convert.channels,
                        ffmpeg_path=config.convert.ffmpeg_path,
                    )
                audio_path = target
                log("  ✓ convert.wav")

            if config.separate.enabled:
                _check_cancel(stop_event)
                log("  [2] Отделение вокала...")
                target = work_dir / "vocals.wav"
                with tempfile.TemporaryDirectory() as tmp:
                    raw = separate_vocals(
                        audio_path, Path(tmp),
                        model=config.separate.model,
                        isolate_stem=config.separate.isolate_stem,
                        unblend_path=config.separate.unblend_path,
                        custom_kernels=config.separate.custom_kernels,
                    )
                    shutil.move(str(raw), str(target))
                audio_path = target
                log("  ✓ vocals.wav")

            if config.transcribe.enabled:
                _check_cancel(stop_event)
                log("  [3] Транскрибация...")
                text = transcribe(
                    audio_path,
                    engine=config.transcribe.engine,
                    model=config.transcribe.model,
                    language=config.transcribe.language,
                    device=config.transcribe.device,
                    compute_type=config.transcribe.compute_type,
                    batch_size=config.transcribe.batch_size,
                    gigaam_provider=config.transcribe.gigaam_provider,
                    local_files_only=config.transcribe.local_files_only,
                    vad_filter=config.transcribe.vad_filter,
                )
                (work_dir / "text.txt").write_text(text, encoding="utf-8")
                log(f"  ✓ text.txt ({len(text)} символов)")

            processed_ok = True

        except Cancelled:
            log("⏹ Отменено")
            break
        except Exception as e:  # noqa: BLE001
            error = f"{type(e).__name__}: {e}"
            log(f"  ✗ Ошибка: {error}")
        finally:
            if processed_ok and not config.keep_intermediate and final_stage is not None:
                try:
                    _persist_final(work_dir, persist_dir, final_stage, log)
                except Exception as e:  # noqa: BLE001
                    error = f"persist failed: {e}"
                    log(f"  ✗ {error}")
            shutil.rmtree(tmp_root, ignore_errors=True)
            if on_progress:
                on_progress(idx, total, "file")
            results.append(FileResult(f, text, persist_dir if not error else None, error))

    return results


# ---------------------------------------------------------------- staged

def _run_staged(
    config: PipelineConfig,
    files: list[Path],
    log: LogFn,
    on_progress: Optional[ProgressFn],
    stop_event: Optional[threading.Event],
) -> list[FileResult]:
    output_root = Path(config.output_dir)
    total = len(files)
    final_stage = _final_stage(config)

    # One tmp_root for the whole batch: stage 2 needs convert.wav of every file,
    # stage 3 needs vocals.wav of every file. All of it is removed at the end.
    tmp_root_ctx = tempfile.mkdtemp(prefix="at_staged_")
    tmp_root = Path(tmp_root_ctx)

    state: dict[Path, dict] = {}
    for f in files:
        work_dir, persist_dir = _work_dir_for(config, f.stem, output_root, tmp_root)
        state[f] = {
            "audio_path": f,
            "work_dir": work_dir,
            "persist_dir": persist_dir,
            "text": None,
            "error": None,
        }

    def active() -> list[Path]:
        return [f for f in files if state[f]["error"] is None]

    def cancelled() -> bool:
        return stop_event is not None and stop_event.is_set()

    try:
        # ---- Stage 1: convert ----
        if config.convert.enabled:
            log(f"\n══ Этап 1/3 · Конвертация ({total} файлов) ══")
            for i, f in enumerate(files, 1):
                if cancelled():
                    log("⏹ Отменено")
                    break
                if on_progress:
                    on_progress(i, total, "convert")
                try:
                    target = state[f]["work_dir"] / "convert.wav"
                    if config.overwrite or not target.exists():
                        convert_to_wav(
                            f, target,
                            sample_rate=config.convert.sample_rate,
                            channels=config.convert.channels,
                            ffmpeg_path=config.convert.ffmpeg_path,
                        )
                    state[f]["audio_path"] = target
                    log(f"  ✓ {f.name}")
                except Exception as e:  # noqa: BLE001
                    state[f]["error"] = f"{type(e).__name__}: {e}"
                    log(f"  ✗ {f.name}: {e}")

        # ---- Stage 2: separate ----
        if config.separate.enabled and not cancelled():
            todo = active()
            log(f"\n══ Этап 2/3 · Отделение вокала ({len(todo)} файлов) ══")
            for i, f in enumerate(todo, 1):
                if cancelled():
                    log("⏹ Отменено")
                    break
                if on_progress:
                    on_progress(i, len(todo), "separate")
                try:
                    target = state[f]["work_dir"] / "vocals.wav"
                    with tempfile.TemporaryDirectory() as tmp:
                        raw = separate_vocals(
                            state[f]["audio_path"], Path(tmp),
                            model=config.separate.model,
                            isolate_stem=config.separate.isolate_stem,
                            unblend_path=config.separate.unblend_path,
                            custom_kernels=config.separate.custom_kernels,
                        )
                        shutil.move(str(raw), str(target))
                    state[f]["audio_path"] = target
                    log(f"  ✓ {f.name}")
                except Exception as e:  # noqa: BLE001
                    state[f]["error"] = f"{type(e).__name__}: {e}"
                    log(f"  ✗ {f.name}: {e}")

        # ---- Stage 3: transcribe ----
        if config.transcribe.enabled and not cancelled():
            todo = active()
            log(f"\n══ Этап 3/3 · Транскрибация ({len(todo)} файлов) ══")
            for i, f in enumerate(todo, 1):
                if cancelled():
                    log("⏹ Отменено")
                    break
                if on_progress:
                    on_progress(i, len(todo), "transcribe")
                try:
                    text = transcribe(
                        state[f]["audio_path"],
                        engine=config.transcribe.engine,
                        model=config.transcribe.model,
                        language=config.transcribe.language,
                        device=config.transcribe.device,
                        compute_type=config.transcribe.compute_type,
                        batch_size=config.transcribe.batch_size,
                        gigaam_provider=config.transcribe.gigaam_provider,
                        local_files_only=config.transcribe.local_files_only,
                        vad_filter=config.transcribe.vad_filter,
                    )
                    (state[f]["work_dir"] / "text.txt").write_text(text, encoding="utf-8")
                    state[f]["text"] = text
                    log(f"  ✓ {f.name} ({len(text)} символов)")
                except Exception as e:  # noqa: BLE001
                    state[f]["error"] = f"{type(e).__name__}: {e}"
                    log(f"  ✗ {f.name}: {e}")

        # ---- Persist final artifacts ----
        if not config.keep_intermediate and final_stage is not None:
            log(f"\n══ Перенос финальных файлов ({final_stage}) ══")
            for f in files:
                s = state[f]
                if s["error"] is not None:
                    continue
                try:
                    _persist_final(s["work_dir"], s["persist_dir"], final_stage, log)
                except Exception as e:  # noqa: BLE001
                    s["error"] = f"persist failed: {e}"
                    log(f"  ✗ {f.name}: {e}")

    finally:
        shutil.rmtree(tmp_root, ignore_errors=True)

    return [
        FileResult(
            f,
            state[f]["text"],
            state[f]["persist_dir"] if state[f]["error"] is None else None,
            state[f]["error"],
        )
        for f in files
    ]