# voice_input.py - macOS Speech framework integration for voice-to-text in the chat app.
# Requires: pyobjc-framework-Speech, pyobjc-framework-AVFoundation
# Usage: request_authorization() then use VoiceRecognitionSession to start/stop recognition.

from __future__ import annotations

import threading
from typing import Callable

# Lazy imports so app runs without Speech/AVFoundation
Speech = None
AVFoundation = None


def _ensure_speech():
    global Speech, AVFoundation
    if Speech is not None:
        return True
    try:
        import Speech as _Speech
        import AVFoundation as _AVF
        Speech = _Speech
        AVFoundation = _AVF
        return True
    except Exception:
        return False


def _has_speech_usage_description():
    """True if the main bundle's Info.plist has NSSpeechRecognitionUsageDescription.
    When running as `python app_pyobjc2.py`, the main bundle is Python's, so this is False and we must not use Speech (TCC would crash)."""
    try:
        from Foundation import NSBundle
        info = NSBundle.mainBundle().infoDictionary()
        if info is None:
            return False
        return bool(info.get("NSSpeechRecognitionUsageDescription"))
    except Exception:
        return False


def is_available():
    """Return True if Speech and AVFoundation are available."""
    return _ensure_speech()


_authorization_status = None  # None = not yet requested, (True, None) = approved, (False, "msg") = denied


def request_authorization_async(callback: Callable[[bool, str | None], None]):
    """Request speech recognition authorization. callback(approved, error_message) is run on main thread when done."""
    global _authorization_status
    if not _has_speech_usage_description():
        callback(
            False,
            "Voice input requires running from an app bundle with NSSpeechRecognitionUsageDescription in Info.plist. "
            "Use the project Info.plist when building with py2app, or run the built .app instead of python app_pyobjc2.py.",
        )
        return
    if not _ensure_speech():
        callback(False, "Speech framework not available.")
        return
    try:
        from PyObjCTools import AppHelper
        def handler(status):
            approved = status == 3
            msg = None
            if status == 1:
                msg = "Speech recognition was denied. Enable it in System Settings > Privacy & Security > Speech Recognition."
            elif status == 2:
                msg = "Speech recognition is restricted on this device."
            def on_main():
                globals()["_authorization_status"] = (approved, msg)
                callback(approved, msg)
            AppHelper.callAfter(on_main)
        Speech.SFSpeechRecognizer.requestAuthorization_(handler)
    except Exception as e:
        try:
            from PyObjCTools import AppHelper
            AppHelper.callAfter(lambda: callback(False, str(e)))
        except Exception:
            callback(False, str(e))


def get_authorization_status():
    """Return (approved: bool, error_message: str|None). Only valid after request_authorization_async has run."""
    return _authorization_status


def supports_on_device_recognition():
    """Return True if the default locale supports on-device (offline) speech recognition.
    Requires dictation enabled and language data downloaded in System Settings."""
    if not _ensure_speech():
        return False
    try:
        recognizer = Speech.SFSpeechRecognizer.alloc().init()
        if recognizer is None or not recognizer.isAvailable():
            return False
        return bool(recognizer.supportsOnDeviceRecognition())
    except Exception:
        return False


class VoiceRecognitionSession:
    """One-shot or toggle session: start() begins listening, stop() ends and returns final text. Partial results via callbacks."""

    def __init__(
        self,
        on_partial: Callable[[str], None] | None = None,
        on_final: Callable[[str], None] | None = None,
        on_error: Callable[[str], None] | None = None,
        on_state: Callable[[str], None] | None = None,
    ):
        self.on_partial = on_partial
        self.on_final = on_final
        self.on_error = on_error
        self.on_state = on_state
        self._engine = None
        self._request = None
        self._task = None
        self._recognizer = None
        self._lock = threading.Lock()
        self._last_transcription = ""
        self._final_result = None

    def _invoke_main(self, fn, *args):
        if fn is None:
            return
        try:
            from PyObjCTools import AppHelper
            AppHelper.callAfter(lambda: fn(*args))
        except Exception:
            try:
                fn(*args)
            except Exception:
                pass

    def start(self, prefer_on_device: bool = False) -> bool:
        """Start listening. If prefer_on_device is True and the system supports it, use on-device (offline) recognition.
        Returns False if framework unavailable or not authorized."""
        with self._lock:
            if self._engine is not None:
                return True  # already running
        if not _has_speech_usage_description():
            self._invoke_main(
                self.on_error,
                "Voice input requires running from an app bundle (Info.plist must include NSSpeechRecognitionUsageDescription).",
            )
            return False
        if not _ensure_speech():
            self._invoke_main(self.on_error, "Speech framework not available.")
            return False
        try:
            recognizer = Speech.SFSpeechRecognizer.alloc().init()
            if recognizer is None or not recognizer.isAvailable():
                self._invoke_main(self.on_error, "Speech recognition is not available on this device.")
                return False
        except Exception as e:
            self._invoke_main(self.on_error, str(e))
            return False

        def do_start():
            with self._lock:
                if self._engine is not None:
                    return
                try:
                    engine = AVFoundation.AVAudioEngine.alloc().init()
                    input_node = engine.inputNode()
                    if input_node is None:
                        self._invoke_main(self.on_error, "No audio input node.")
                        return
                    format = input_node.outputFormatForBus_(0)
                    if format is None:
                        self._invoke_main(self.on_error, "Could not get input format.")
                        return
                    req = Speech.SFSpeechAudioBufferRecognitionRequest.alloc().init()
                    req.setShouldReportPartialResults_(True)
                    try:
                        req.setAddsPunctuation_(True)
                    except Exception:
                        pass
                    # Prefer on-device (offline) recognition when requested and supported
                    if prefer_on_device and getattr(recognizer, "supportsOnDeviceRecognition", None) and recognizer.supportsOnDeviceRecognition():
                        try:
                            req.setRequiresOnDeviceRecognition_(True)
                        except Exception:
                            pass
                    self._request = req
                    self._engine = engine
                    self._recognizer = recognizer
                    self._last_transcription = ""
                    self._final_result = None

                    def buffer_callback(buffer, _):
                        if self._request is not None:
                            self._request.appendAudioPCMBuffer_(buffer)

                    input_node.installTapOnBus_bufferSize_format_block_(
                        0, 4096, format, buffer_callback
                    )
                    engine.prepare()
                    result = engine.startAndReturnError_(None)
                    if isinstance(result, tuple):
                        success, err_obj = result[0], result[1] if len(result) > 1 else None
                    else:
                        success, err_obj = bool(result), None
                    if not success:
                        msg = str(err_obj) if err_obj else "Failed to start audio engine."
                        self._invoke_main(self.on_error, msg)
                        self._cleanup()
                        return
                    self._invoke_main(self.on_state, "listening")

                    def _normalize_speech_error(err_str):
                        s = (err_str or "").lower()
                        if "canceled" in s or "cancelled" in s:
                            return None
                        if "retry" in s or "203" in err_str or "kafassistanterrordomain" in s or "sirispeecherrordomain" in s:
                            return (
                                "Speech recognition is temporarily unavailable. "
                                "Check your internet connection (recognition uses Apple’s servers) and try again."
                            )
                        return err_str

                    def result_handler(result, error):
                        if error is not None:
                            err_str = str(error)
                            if "canceled" in err_str.lower() or "cancelled" in err_str.lower():
                                return
                            friendly = _normalize_speech_error(err_str)
                            self._invoke_main(self.on_error, friendly or err_str)
                            return
                        if result is None:
                            return
                        try:
                            best = result.bestTranscription()
                            if best is None:
                                return
                            text = best.formattedString()
                            if not text:
                                return
                            try:
                                is_final = result.isFinal()
                            except Exception:
                                is_final = False
                            with self._lock:
                                self._last_transcription = text
                                if is_final:
                                    self._final_result = text
                            if is_final:
                                self._invoke_main(self.on_final, text)
                            else:
                                self._invoke_main(self.on_partial, text)
                        except Exception as e:
                            self._invoke_main(self.on_error, str(e))

                    self._task = recognizer.recognitionTaskWithRequest_resultHandler_(req, result_handler)
                except Exception as e:
                    self._invoke_main(self.on_error, str(e))
                    self._cleanup()

        thread = threading.Thread(target=do_start, daemon=True)
        thread.start()
        return True

    def _cleanup(self):
        with self._lock:
            req = self._request
            engine = self._engine
            task = self._task
            input_node = engine.inputNode() if engine else None
            self._request = None
            self._engine = None
            self._task = None
        try:
            if task is not None:
                try:
                    task.cancel()
                except Exception:
                    pass
            if engine is not None:
                if input_node is not None:
                    try:
                        input_node.removeTapOnBus_(0)
                    except Exception:
                        pass
                try:
                    engine.stop()
                except Exception:
                    pass
        except Exception:
            pass

    def stop(self) -> str | None:
        """Stop listening and return the final transcription (or last partial if no final)."""
        with self._lock:
            req = self._request
            engine = self._engine
            task = self._task
            input_node = self._engine.inputNode() if self._engine else None
        if engine is None:
            with self._lock:
                return self._final_result or self._last_transcription or None
        try:
            if req is not None:
                try:
                    req.endAudio()
                except Exception:
                    pass
            # Give the recognition task a moment to deliver final result
            import time
            time.sleep(0.35)
        except Exception:
            pass
        with self._lock:
            final = self._final_result
            last = self._last_transcription
            self._request = None
            self._engine = None
            self._task = None
        try:
            if task is not None:
                try:
                    task.cancel()
                except Exception:
                    pass
            if input_node is not None:
                try:
                    input_node.removeTapOnBus_(0)
                except Exception:
                    pass
            engine.stop()
        except Exception:
            pass
        return final if final else (last if last else None)

    def is_active(self) -> bool:
        with self._lock:
            return self._engine is not None
