from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
import argparse
import atexit
import json
import os
import re
import signal
import subprocess
import tempfile
import threading
from datetime import datetime

HOST = "127.0.0.1"
VERSION = "3.1"
APP_NAME = "AI Voice Bridge"
SUPPORTED_ORIGINS = {
    "https://chatgpt.com",
    "https://chat.openai.com",
    "https://claude.ai",
    "https://chat.deepseek.com",
    "https://gemini.google.com",
    "https://grok.com",
    "https://www.grok.com",
    "https://grok.x.ai",
}
MAX_TRANSCRIBE_BYTES = 30 * 1024 * 1024

current_proc = None
stop_generation = 0
proc_lock = threading.Lock()
server_config = {}
_whisper_model_cache = None
_whisper_model_lock = threading.Lock()
_shutdown_cleanup_lock = threading.Lock()
_shutdown_cleanup_done = False


def timestamp():
    return datetime.now().strftime("%H:%M:%S")


def _terminate_proc(proc):
    if not proc or proc.poll() is not None:
        return

    try:
        proc.terminate()
        proc.wait(timeout=0.5)
        return
    except Exception:
        pass

    try:
        subprocess.run(
            ["taskkill", "/PID", str(proc.pid), "/T", "/F"],
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
            timeout=1.5,
        )
    except Exception:
        pass


def _stop_voice_locked():
    global current_proc

    proc = current_proc
    current_proc = None
    _terminate_proc(proc)


def stop_voice():
    global stop_generation

    with proc_lock:
        stop_generation += 1
        _stop_voice_locked()


def _cleanup_on_exit():
    """Stop native TTS exactly once when the local server exits normally."""
    global _shutdown_cleanup_done

    with _shutdown_cleanup_lock:
        if _shutdown_cleanup_done:
            return
        _shutdown_cleanup_done = True

    stop_voice()


def _interrupt_for_shutdown(signum, frame):
    # Convert normal termination signals into the same graceful path as Ctrl+C.
    raise KeyboardInterrupt


def clean_tts_text(text):
    text = (text or "").strip()
    if not text:
        return ""

    replacements = {
        "\u2018": "'", "\u2019": "'", "\u201a": "'", "\u201b": "'",
        "\u201c": '"', "\u201d": '"', "\u201e": '"', "\u201f": '"',
        "\u2013": "-", "\u2014": "-", "\u2015": "-", "\u2212": "-",
        "\u2026": "...", "\u00a0": " ",
    }
    for bad, good in replacements.items():
        text = text.replace(bad, good)

    text = re.sub(r"\bLET'?S GO\b", "Let's go", text)
    text = re.sub(r"\bLets go\b", "Let's go", text, flags=re.IGNORECASE)
    text = re.sub(r"\bv(\d+)\.(\d+)\b", r"version \1 point \2", text, flags=re.IGNORECASE)
    text = re.sub(r"[ \t]{2,}", " ", text)
    return text.strip()


def speak_text(text, voice_name, rate, volume, worker_path):
    global current_proc

    text = clean_tts_text(text)
    if not text:
        return

    with proc_lock:
        request_stop_generation = stop_generation

    fd, text_path = tempfile.mkstemp(prefix="rits_voice_bridge_", suffix=".txt", text=True)
    with os.fdopen(fd, "w", encoding="utf-8") as f:
        f.write(text)

    cmd = [
        "powershell.exe",
        "-NoProfile",
        "-ExecutionPolicy", "Bypass",
        "-File", worker_path,
        "-TextFile", text_path,
        "-VoiceName", voice_name,
        "-Rate", str(rate),
        "-Volume", str(volume),
    ]

    with proc_lock:
        if stop_generation != request_stop_generation:
            print("SPEAK CANCELLED: stop requested before TTS process started")
            return

        _stop_voice_locked()
        current_proc = subprocess.Popen(
            cmd,
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
        )



def is_whisper_module_available():
    try:
        import faster_whisper  # noqa: F401
        return True
    except Exception:
        return False


def transcribe_whisper(wav_path):
    global _whisper_model_cache

    model_name = server_config.get("whisper_model", "base.en")
    model_root = server_config.get("whisper_model_root", "")
    device = server_config.get("whisper_device", "cpu")
    compute_type = server_config.get("whisper_compute_type", "int8")

    try:
        from faster_whisper import WhisperModel
    except Exception as e:
        return False, "", f"Whisper Python package is not installed or failed to load: {e}. Run Setup-Local-Dictation-Windows.bat."

    try:
        with _whisper_model_lock:
            if _whisper_model_cache is None:
                print(f"[{timestamp()}] Loading Whisper model: {model_name}")
                _whisper_model_cache = WhisperModel(
                    model_name,
                    device=device,
                    compute_type=compute_type,
                    download_root=model_root or None,
                )
            model = _whisper_model_cache

        segments, info = model.transcribe(
            wav_path,
            language="en",
            beam_size=5,
            vad_filter=True,
        )
        text = " ".join((segment.text or "").strip() for segment in segments).strip()
        return True, text, ""

    except Exception as e:
        return False, "", f"Whisper transcription failed: {e}"


def make_handler(voice_name, rate, volume, worker_path):
    class VoiceBridgeHandler(BaseHTTPRequestHandler):
        def _send_cors(self):
            origin = self.headers.get("Origin", "")
            if origin in SUPPORTED_ORIGINS or origin.startswith("chrome-extension://"):
                self.send_header("Access-Control-Allow-Origin", origin)
                self.send_header("Vary", "Origin")
            elif not origin:
                self.send_header("Access-Control-Allow-Origin", "*")
            self.send_header("Access-Control-Allow-Methods", "POST, GET, OPTIONS")
            self.send_header("Access-Control-Allow-Headers", "Content-Type")
            self.send_header("Access-Control-Allow-Private-Network", "true")

        def _json(self, status, payload):
            self.send_response(status)
            self._send_cors()
            self.send_header("Content-Type", "application/json; charset=utf-8")
            self.end_headers()
            self.wfile.write(json.dumps(payload).encode("utf-8"))

        def do_OPTIONS(self):
            self.send_response(204)
            self._send_cors()
            self.end_headers()

        def do_GET(self):
            if self.path == "/health":
                self._json(200, {
                    "ok": True,
                    "service": APP_NAME,
                    "version": VERSION,
                    "tts": True,
                    "dictation": bool(server_config.get("whisper_available")),
                    "stt_backend": "whisper" if server_config.get("whisper_available") else "",
                    "whisper_model": server_config.get("whisper_model", ""),
                    "whisper_device": server_config.get("whisper_device", "cpu"),
                    "whisper_compute_type": server_config.get("whisper_compute_type", "int8"),
                })
                return
            self._json(404, {"ok": False, "error": "Not found"})

        def do_POST(self):
            print(f"[{timestamp()}] POST {self.path}")

            if self.path == "/stop":
                stop_voice()
                print("VOICE STOPPED")
                self._json(200, {"ok": True, "stopped": True})
                return

            if self.path == "/speak":
                self._handle_speak()
                return

            if self.path == "/transcribe":
                self._handle_transcribe()
                return

            self._json(404, {"ok": False, "error": "Not found"})

        def _handle_speak(self):
            length = int(self.headers.get("Content-Length", 0))
            body = self.rfile.read(length)

            try:
                data = json.loads(body.decode("utf-8"))
                text = data.get("text", "").strip()
            except Exception as e:
                print("Bad JSON:", e)
                text = ""

            print("TEXT RECEIVED:", repr(text[:200]))

            if text:
                speak_text(text, voice_name, rate, volume, worker_path)

            self._json(200, {"ok": True})

        def _handle_transcribe(self):
            if not server_config.get("whisper_available"):
                self._json(503, {"ok": False, "error": "Local dictation is not available. Run Setup-Local-Dictation-Windows.bat, then restart this server."})
                return

            length = int(self.headers.get("Content-Length", 0))
            if length <= 0:
                self._json(400, {"ok": False, "error": "No audio data received."})
                return

            if length > MAX_TRANSCRIBE_BYTES:
                self._json(413, {"ok": False, "error": "Recording is too large. Try a shorter dictation."})
                return

            audio = self.rfile.read(length)
            if not audio:
                self._json(400, {"ok": False, "error": "No audio data received."})
                return

            fd, wav_path = tempfile.mkstemp(prefix="rits_dictation_", suffix=".wav")
            try:
                with os.fdopen(fd, "wb") as f:
                    f.write(audio)

                print(f"[{timestamp()}] WAV RECEIVED: {len(audio)} bytes")
                print(f"[{timestamp()}] TRANSCRIBE running Whisper local backend")
                ok, text, error = transcribe_whisper(wav_path)

                if ok:
                    print("TRANSCRIPT:", repr(text[:200]))
                    self._json(200, {"ok": True, "text": text, "backend": "whisper"})
                else:
                    print("TRANSCRIBE ERROR:", error)
                    self._json(500, {"ok": False, "error": error, "backend": "whisper"})
            finally:
                try:
                    os.remove(wav_path)
                except Exception:
                    pass

        def log_message(self, format, *args):
            return

    return VoiceBridgeHandler


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=f"{APP_NAME} v{VERSION}")
    parser.add_argument("--port", type=int, default=5055)
    parser.add_argument("--voice", default="")
    parser.add_argument("--rate", type=int, default=0)
    parser.add_argument("--volume", type=int, default=100)
    parser.add_argument("--whisper-model", default="base.en")
    parser.add_argument("--whisper-device", default="cpu")
    parser.add_argument("--whisper-compute-type", default="int8")
    args = parser.parse_args()

    base_dir = os.path.dirname(os.path.abspath(__file__))
    worker_path = os.path.join(base_dir, "speech_worker.ps1")

    server_config["whisper_model"] = args.whisper_model
    server_config["whisper_model_root"] = os.path.join(base_dir, "whisper-models")
    server_config["whisper_device"] = args.whisper_device
    server_config["whisper_compute_type"] = args.whisper_compute_type
    server_config["whisper_available"] = is_whisper_module_available()

    print("=" * 66)
    print(f"{APP_NAME} {VERSION}")
    print(f"Server: http://{HOST}:{args.port}")
    print(f"Voice: {args.voice or 'Windows default voice'}")
    print(f"Rate: {args.rate}")
    print(f"Volume: {args.volume}")
    print("TTS output: available")

    if server_config["whisper_available"]:
        print("Local dictation: available (Whisper accuracy backend)")
        print(f"Whisper model: {server_config['whisper_model']}")
        print(f"Whisper device: {server_config['whisper_device']}")
        print(f"Whisper compute: {server_config['whisper_compute_type']}")
    else:
        print("Local dictation: not available yet")
        print("To enable: run Setup-Local-Dictation-Windows.bat, then restart this server")

    print("=" * 66)
    print("Waiting for requests...")
    print("Press Ctrl+C in this window to stop the server.")
    print()

    handler = make_handler(
        voice_name=args.voice,
        rate=args.rate,
        volume=args.volume,
        worker_path=worker_path,
    )

    atexit.register(_cleanup_on_exit)
    for signal_name in ("SIGTERM", "SIGHUP"):
        shutdown_signal = getattr(signal, signal_name, None)
        if shutdown_signal is not None:
            try:
                signal.signal(shutdown_signal, _interrupt_for_shutdown)
            except (OSError, RuntimeError, ValueError):
                pass

    httpd = ThreadingHTTPServer((HOST, args.port), handler)
    try:
        httpd.serve_forever()
    except KeyboardInterrupt:
        print(f"\n{APP_NAME} stopped.")
    finally:
        _cleanup_on_exit()
        httpd.server_close()
