"""Audio captcha solver.

Strategy:
  1. faster-whisper (tiny) with beam_size=1, language='en'
  2. If ollama enabled, ask it to clean the digits
  3. Local LLM (Qwen2.5-1.5B) as a digit-extraction fallback

We accept:
  - MP3, WAV, OGG (any ffmpeg-decodable audio)
  - Tries to extract ONLY digits if the text is clearly digit-only audio
  - Returns the full transcript otherwise
"""

from __future__ import annotations

import re
from typing import Optional

from captcha_solver.solvers.base import BaseSolver, SolveAttempt
from captcha_solver.utils.image import decode_base64_audio


class AudioSolver(BaseSolver):
    name = "audio"
    captcha_type = "audio"

    def __init__(self, ctx) -> None:
        super().__init__(ctx)
        self._audio: bytes = b""

    def prepare(self, image_b64: Optional[str], audio_b64: Optional[str], hint: Optional[str]) -> None:
        if not audio_b64:
            self._audio = b""
            return
        try:
            self._audio = decode_base64_audio(audio_b64)
        except Exception as exc:
            self._audio = b""
            self._last_error = f"decode: {exc}"

    def attempts(self):
        return [
            self._whisper_only,
            self._whisper_plus_llm,
        ]

    def _whisper_only(self) -> SolveAttempt:
        if not self._audio:
            return SolveAttempt(answer="", confidence=0.0, solver_name="audio.whisper", error="no audio")
        text, conf = self.ctx.whisper.transcribe(self._audio, language="en", beam_size=1)
        digits = _extract_digits(text)
        if digits and len(digits) >= 3:
            return SolveAttempt(answer=digits, confidence=max(conf, 0.6), solver_name="audio.whisper.digits")
        if text:
            return SolveAttempt(answer=text, confidence=conf, solver_name="audio.whisper.text")
        return SolveAttempt(answer="", confidence=0.0, solver_name="audio.whisper", error="empty transcript")

    def _whisper_plus_llm(self) -> SolveAttempt:
        if not self._audio:
            return SolveAttempt(answer="", confidence=0.0, solver_name="audio.whisper_llm", error="no audio")
        text, conf = self.ctx.whisper.transcribe(self._audio, language="en", beam_size=5)
        if not text:
            return SolveAttempt(answer="", confidence=0.0, solver_name="audio.whisper_llm", error="empty")
        try:
            cleaned = self.ctx.qwen.generate(
                f"Extract only the digits/words from this audio transcript. "
                f"Return them in the order spoken, space-separated. No commentary.\n\nTranscript: {text}",
                max_new_tokens=40,
            )
        except Exception as exc:
            return SolveAttempt(answer="", confidence=0.0, solver_name="audio.whisper_llm", error=str(exc))
        cleaned = cleaned.strip()
        if not cleaned:
            return SolveAttempt(answer="", confidence=0.0, solver_name="audio.whisper_llm", error="llm empty")
        return SolveAttempt(answer=cleaned, confidence=min(conf + 0.1, 0.95), solver_name="audio.whisper_llm")


def _extract_digits(text: str) -> str:
    digits = re.findall(r"\d", text)
    if len(digits) >= 3:
        return " ".join(digits)
    return ""