from __future__ import annotations

import json
import logging
import time

from app import diarize, media, store, transcribe, youtube
from app.config import get_settings
from app.eligibility import assess
from app.errors import JobCancelled, VaaniError
from app.languages import display_name, script_name, script_warning
from app.codeswitch import blend_words
from app.transcript import merge_turns, project_speakers, render
from app.util import free_gb, utcnow

log = logging.getLogger("vaani.pipeline")


def run(job_id: str) -> None:
    current = {"name": "preparing"}
    try:
        _ensure_disk()
        if _cancelled(job_id):
            raise JobCancelled()
        job = store.load(job_id)
        if job["steps"]["preparing"]["state"] != "done":
            current["name"] = "preparing"
            _prepare(job_id)
        if _cancelled(job_id):
            raise JobCancelled()
        job = store.load(job_id)
        if job["steps"]["audio"]["state"] != "done":
            current["name"] = "audio"
            _audio(job_id)
        if _cancelled(job_id):
            raise JobCancelled()
        if store.load(job_id)["steps"]["transcript"]["state"] != "done":
            current["name"] = "transcript"
            _transcript(job_id)

        def finish(current_job: dict) -> None:
            current_job["status"] = "complete"
            current_job["error"] = None
            current_job["steps"]["ready"]["state"] = "done"
            current_job["steps"]["ready"]["detail"] = "Ready"
            current_job["steps"]["ready"]["ended_at"] = utcnow()
            current_job["progress"]["detail"] = "Ready"

        store.update(job_id, finish)
        log.info("job %s complete", job_id)
    except JobCancelled:
        _fail(job_id, current["name"], "cancelled", "Stopped.", "retry")
    except VaaniError as exc:
        log.warning("job %s failed: %s", job_id, exc.message)
        _fail(job_id, current["name"], exc.code, exc.message, exc.recovery)
    except Exception as exc:
        log.exception("job %s crashed", job_id)
        message = f"The job stopped because of a problem on this computer. {type(exc).__name__}: {exc}"
        _fail(job_id, current["name"], "unexpected", message[:300], "retry")


def _prepare(job_id: str) -> None:
    _set(job_id, "preparing", "active", "Reading the video")
    job = store.load(job_id)
    source = job["source"]
    settings = get_settings()
    directory = store.job_dir(job_id)
    if source.get("type") == "upload" or media.find_named(directory, "upload"):
        path = media.find_named(directory, "upload")
        if path is None:
            raise VaaniError(
                "The uploaded file is missing.",
                "audio_unreadable",
                400,
                "upload",
            )
        duration = source.get("duration")
        if not duration:
            duration = int(round(media.probe_duration(path)))
        allowed, code, message = assess(
            {"duration": duration, "live_status": "not_live"},
            settings.max_duration_seconds,
        )
        if not allowed:
            raise VaaniError(message or "This file cannot be used.", code or "too_long", 400)

        def mutate(current_job: dict, duration=duration) -> None:
            current_job["source"]["duration"] = duration
            current_job["source"]["title"] = current_job["source"].get("title") or "Uploaded recording"

        store.update(job_id, mutate)
        _set(job_id, "preparing", "done", "No thumbnail for an uploaded file")
        return

    info = youtube.fetch_video(source["url"])
    allowed, code, message = assess(info, settings.max_duration_seconds)
    if not allowed:
        raise VaaniError(message or "This video cannot be used.", code or "not_eligible", 400)
    thumb = directory / "thumbnail.jpg"
    saved = youtube.save_thumbnail(info["video_id"], thumb)

    def mutate(current_job: dict, info=info, saved=saved) -> None:
        current_job["source"].update(
            {
                "title": info["title"],
                "channel": info["channel"],
                "duration": info["duration"],
                "upload_date": info["upload_date"],
                "video_id": info["video_id"],
                "webpage_url": info["webpage_url"],
            }
        )
        if saved:
            current_job["artifacts"]["thumbnail"] = "thumbnail.jpg"
        elif "The thumbnail was not available." not in current_job["warnings"]:
            current_job["warnings"].append("The thumbnail was not available.")

    store.update(job_id, mutate)
    _set(job_id, "preparing", "done", "Thumbnail saved" if saved else "The thumbnail was not available")


def _audio(job_id: str) -> None:
    directory = store.job_dir(job_id)
    mp3 = directory / "audio.mp3"
    if mp3.exists() and mp3.stat().st_size > 1000:
        def done(current_job: dict) -> None:
            current_job["steps"]["audio"]["state"] = "done"
            current_job["steps"]["audio"]["detail"] = "MP3 ready"
            current_job["steps"]["audio"]["ended_at"] = current_job["steps"]["audio"]["ended_at"] or utcnow()
            current_job["artifacts"]["audio"] = "audio.mp3"

        store.update(job_id, done)
        return

    _set(job_id, "audio", "active", "Writing the MP3")
    job = store.load(job_id)
    source_file = media.find_named(directory, "upload") or media.find_named(directory, "source")
    if source_file is None:
        if job["source"].get("type") != "youtube":
            raise VaaniError("The uploaded file is missing.", "audio_unreadable", 400, "upload")
        _set(job_id, "audio", "active", "Downloading audio")
        source_file = youtube.download_audio(
            job["source"]["url"],
            directory,
            lambda: _cancelled(job_id),
            lambda detail: _set(job_id, "audio", "active", detail),
        )
    _set(job_id, "audio", "active", "Writing the MP3")
    media.render_audio(
        source_file,
        mp3,
        directory / "audio.wav",
        job["source"].get("title") or "Vaani",
        job["source"].get("channel") or "",
        lambda: _cancelled(job_id),
    )
    if source_file.name.startswith("source."):
        source_file.unlink(missing_ok=True)

    def ready(current_job: dict) -> None:
        current_job["artifacts"]["audio"] = "audio.mp3"

    store.update(job_id, ready)
    _set(job_id, "audio", "done", "MP3 ready")


def _transcript(job_id: str) -> None:
    _set(
        job_id,
        "transcript",
        "active",
        "Loading the speech model. The first time, it downloads and can take several minutes.",
    )
    directory = store.job_dir(job_id)
    mp3 = directory / "audio.mp3"
    wav = directory / "audio.wav"
    if not wav.exists() or wav.stat().st_size < 1000:
        if not mp3.exists():
            raise VaaniError(
                "The MP3 is missing, so there is nothing to transcribe.",
                "audio_failed",
                500,
                "retry",
            )
        media.wav_from_mp3(mp3, wav, lambda: _cancelled(job_id))
    job = store.load(job_id)
    language = job["options"].get("language") or "auto"
    requested = None if language == "auto" else language
    last_note = {"at": 0.0}
    prep_notes: list[str] = []
    listen = wav
    if mp3.exists():
        _set(job_id, "transcript", "active", "Removing background music")
        speech = directory / "speech.wav"
        try:
            media.isolate_speech(mp3, speech, lambda: _cancelled(job_id))
            listen = speech
        except JobCancelled:
            raise
        except VaaniError:
            log.exception("could not pull the music down for %s", job_id)
            prep_notes.append("Background music could not be reduced, so the full mix was used.")

    def on_segment(index: int, detected: str) -> None:
        if _cancelled(job_id):
            raise JobCancelled()
        now = time.monotonic()
        if index > 1 and now - last_note["at"] < 0.8:
            return
        last_note["at"] = now
        name = display_name(detected)
        detail = f"Listening in {name}" if index <= 0 else f"Listening in {name}, segment {index}"
        _set(job_id, "transcript", "active", detail)

    if requested == "en":
        chunks, detected = transcribe.transcribe_wav(
            listen,
            requested,
            on_segment,
            lambda: _cancelled(job_id),
        )
    else:
        words, detected = transcribe.transcribe_words(
            listen,
            requested,
            on_segment,
            lambda: _cancelled(job_id),
        )
        if requested == "hi" or detected == "hi":
            _set(job_id, "transcript", "active", "Checking which words are English")
            english_words, _english_detected = transcribe.transcribe_words(
                listen,
                "en",
                None,
                lambda: _cancelled(job_id),
            )
            chunks = blend_words(words, english_words)
        else:
            chunks = blend_words(words, [])
    samples, rate = diarize.read_wav(listen)
    speakers = int(job["options"].get("speakers") or 2)
    _set(job_id, "transcript", "active", "Separating voices")
    labeled, speaker_warning = diarize.label_speakers(chunks, samples, rate, speakers)
    turns = merge_turns(labeled)
    warnings: list[str] = list(prep_notes)
    if not turns:
        warnings.append("No speech was found in this recording.")
    if speaker_warning:
        warnings.append(speaker_warning)
    spoken_text = render(turns)
    script_note = script_warning(detected, spoken_text)
    if script_note:
        warnings.append(script_note)

    english_turns: list[dict] = []
    if job["options"].get("english") and detected and detected != "en":
        _set(job_id, "transcript", "active", "Writing the English transcript")
        try:
            english_chunks = transcribe.translate_wav(
                listen,
                detected,
                lambda: _cancelled(job_id),
            )
            english_turns = project_speakers(turns, english_chunks)
        except JobCancelled:
            raise
        except VaaniError as exc:
            warnings.append(exc.message)
        except Exception:
            log.exception("english transcript failed for %s", job_id)
            warnings.append("The English transcript could not be written. The spoken transcript is ready.")

    (directory / "transcript.txt").write_text(spoken_text, encoding="utf-8-sig")
    artifacts = {"transcript": "transcript.txt", "turns": "turns.json"}
    english_text = render(english_turns)
    if english_text:
        (directory / "transcript.en.txt").write_text(english_text, encoding="utf-8-sig")
        artifacts["transcript_en"] = "transcript.en.txt"
    payload = {
        "language": detected,
        "language_name": display_name(detected),
        "script": script_name(detected),
        "turns": turns,
        "english_turns": english_turns,
    }
    (directory / "turns.json").write_text(
        json.dumps(payload, ensure_ascii=False, indent=2),
        encoding="utf-8",
    )

    def mutate(current_job: dict, warnings=list(warnings), detected=detected, artifacts=dict(artifacts)) -> None:
        kept = [
            item
            for item in current_job.get("warnings") or []
            if str(item).startswith("The thumbnail")
        ]
        current_job["detected_language"] = detected
        current_job["warnings"] = _unique(kept + warnings)
        current_job["artifacts"].update(artifacts)

    store.update(job_id, mutate)
    wav.unlink(missing_ok=True)
    (directory / "speech.wav").unlink(missing_ok=True)
    _set(job_id, "transcript", "done", "Transcript written")


def _ensure_disk() -> None:
    settings = get_settings()
    if free_gb(settings.data_dir) < settings.min_free_gb:
        raise VaaniError(
            "This computer is low on free disk space, so a new recording cannot be saved.",
            "disk_low",
            400,
        )


def _cancelled(job_id: str) -> bool:
    try:
        return bool(store.load(job_id).get("cancel_requested"))
    except VaaniError:
        return True


def _set(job_id: str, step: str, state: str, detail: str = "") -> None:
    def mutate(current_job: dict) -> None:
        item = current_job["steps"][step]
        if state == "active" and not item["started_at"]:
            item["started_at"] = utcnow()
        if state in {"done", "failed"}:
            item["ended_at"] = utcnow()
        item["state"] = state
        if detail:
            item["detail"] = detail
        if state == "active":
            current_job["status"] = step
            current_job["progress"]["detail"] = detail or step

    store.update(job_id, mutate)


def _fail(job_id: str, step: str, code: str, message: str, recovery: str | None) -> None:
    def mutate(current_job: dict) -> None:
        step_name = step if step in current_job["steps"] else "preparing"
        item = current_job["steps"][step_name]
        item["state"] = "failed"
        item["detail"] = message
        item["ended_at"] = utcnow()
        current_job["status"] = "failed"
        current_job["error"] = {"code": code, "message": message, "recovery": recovery}
        current_job["progress"]["detail"] = message

    try:
        store.update(job_id, mutate)
    except VaaniError:
        log.warning("could not record failure for %s", job_id)


def _unique(items: list[str]) -> list[str]:
    seen: list[str] = []
    for item in items:
        if item and item not in seen:
            seen.append(item)
    return seen