#!/usr/bin/env python3
"""
Jotbot: Push-to-talk voice dictation for macOS.

Uses Deepgram for real-time streaming transcription (preferred) with Whisper as fallback.
Optional Haiku formatting cleans up transcribed text for code-related content.

Hotkeys:
    Ctrl+Option       - Record and transcribe (raw)
    Cmd+Ctrl+Option   - Record and transcribe with Haiku formatting (press Cmd first)
"""

import argparse
import io
import json
import os
import sys
import threading
import time
import wave
from pathlib import Path

import anthropic
import numpy as np
import pyperclip
import sounddevice as sd
from dotenv import load_dotenv
from openai import OpenAI
from pynput.keyboard import Controller, Key, Listener

# Config location: ~/.config/jotbot/.env
CONFIG_DIR = Path.home() / ".config" / "jotbot"
CONFIG_FILE = CONFIG_DIR / ".env"


def load_config():
    """Load environment variables from config file or local .env."""
    # Priority: local .env > ~/.config/jotbot/.env > environment
    if Path(".env").exists():
        load_dotenv(Path(".env"))
    elif CONFIG_FILE.exists():
        load_dotenv(CONFIG_FILE)
    else:
        load_dotenv()  # Just load from environment


load_config()

# Deepgram imports are done lazily in run() after config is loaded

# Audio configuration
SAMPLE_RATE = 16000  # Whisper's native sample rate
CHANNELS = 1
BLOCKSIZE = 4000  # ~250ms chunks at 16kHz, good for streaming

# Hotkey definitions
HOTKEY_BASE = {Key.ctrl, Key.alt}
HOTKEY_FORMAT = {Key.ctrl, Key.alt, Key.cmd}

# Haiku formatting system prompt
HAIKU_SYSTEM = """You are a text formatter. You receive voice-transcribed text and output a cleaned-up version.

Rules:
- Format code identifiers with backticks: function names, variables, file paths, class names
- Fix punctuation and capitalization
- Preserve the meaning exactly - don't add, remove, or interpret anything
- If it's not code-related, just clean it up minimally
- Output ONLY the formatted text, no preamble, no explanation"""


class Recorder:
    """Records audio from the microphone and optionally streams to a transcriber."""

    def __init__(self):
        self.audio_chunks = []
        self.stream = None
        self.recording = False
        self.transcriber = None

    def start(self, transcriber=None) -> bool:
        """Start recording audio. Returns True on success."""
        self.audio_chunks = []
        self.transcriber = transcriber

        try:
            self.stream = sd.InputStream(
                samplerate=SAMPLE_RATE,
                channels=CHANNELS,
                dtype=np.float32,
                blocksize=BLOCKSIZE,
                callback=self._audio_callback,
            )
            self.stream.start()
            self.recording = True
            return True
        except sd.PortAudioError as e:
            print(f"Audio error: {e}")
            self.recording = False
            return False

    def stop(self) -> bytes:
        """Stop recording and return WAV audio data."""
        self.recording = False
        self.transcriber = None

        if self.stream:
            self.stream.stop()
            self.stream.close()
            self.stream = None

        print("Processing...")
        return self._to_wav()

    def _audio_callback(self, indata, frames, time_info, status):
        """Called by sounddevice for each audio chunk."""
        if not self.recording:
            return

        self.audio_chunks.append(indata.copy())

        # Stream to Deepgram in real-time if transcriber is active
        if self.transcriber:
            audio_int16 = (indata * 32767).astype(np.int16)
            self.transcriber.send_audio(audio_int16.tobytes())

    def _to_wav(self) -> bytes:
        """Convert recorded audio chunks to WAV format."""
        if not self.audio_chunks:
            return b""

        audio = np.concatenate(self.audio_chunks)
        audio_int16 = (audio * 32767).astype(np.int16)

        buffer = io.BytesIO()
        with wave.open(buffer, "wb") as wf:
            wf.setnchannels(CHANNELS)
            wf.setsampwidth(2)  # 16-bit
            wf.setframerate(SAMPLE_RATE)
            wf.writeframes(audio_int16.tobytes())

        buffer.seek(0)
        return buffer.read()


class DeepgramTranscriber:
    """Persistent streaming transcription using Deepgram WebSocket API."""

    def __init__(self):
        from deepgram import DeepgramClient

        self.client = DeepgramClient()
        self.socket = None
        self.context_manager = None
        self.transcript_parts = []
        self.lock = threading.Lock()
        self.is_connected = False
        self.is_recording = False
        self.should_stop = False
        self.listen_thread = None
        self.keepalive_thread = None
        self._waiting_for_final = False

    def connect(self) -> bool:
        """Open persistent WebSocket connection to Deepgram."""
        from deepgram.core.events import EventType

        if self.is_connected:
            return True

        try:
            self.should_stop = False
            self.context_manager = self.client.listen.v1.connect(
                model="nova-2",
                encoding="linear16",
                sample_rate=str(SAMPLE_RATE),
                channels="1",
                punctuate="true",
                smart_format="true",
                interim_results="true",
                endpointing="300",  # Detect end of speech after 300ms silence
            )
            self.socket = self.context_manager.__enter__()

            self.socket.on(EventType.OPEN, self._on_open)
            self.socket.on(EventType.MESSAGE, self._on_message)
            self.socket.on(EventType.ERROR, self._on_error)
            self.socket.on(EventType.CLOSE, self._on_close)

            self.listen_thread = threading.Thread(target=self._listen_loop, daemon=True)
            self.listen_thread.start()

            self.keepalive_thread = threading.Thread(target=self._keepalive_loop, daemon=True)
            self.keepalive_thread.start()

            # Wait for connection (up to 2 seconds)
            for _ in range(200):
                if self.is_connected:
                    return True
                time.sleep(0.01)

            return self.is_connected

        except Exception as e:
            print(f"Deepgram connection error: {e}")
            self._cleanup()
            return False

    def start_recording(self):
        """Start a new recording session (resets transcript)."""
        with self.lock:
            self.transcript_parts = []
            self._waiting_for_final = False
        self.is_recording = True

    def stop_recording(self) -> str:
        """Stop recording and return accumulated transcript."""
        # Send Finalize message to flush Deepgram's buffer
        if self.socket and self.is_connected:
            try:
                self.socket._websocket.send(json.dumps({"type": "Finalize"}))
            except Exception:
                pass

        with self.lock:
            self._waiting_for_final = True

        # Wait for final results (up to 1.5 seconds)
        deadline = time.time() + 1.5
        last_part_count = 0
        while time.time() < deadline:
            time.sleep(0.1)
            with self.lock:
                if len(self.transcript_parts) > last_part_count:
                    last_part_count = len(self.transcript_parts)
                    deadline = min(deadline, time.time() + 0.5)

        self.is_recording = False
        with self.lock:
            self._waiting_for_final = False
            return " ".join(self.transcript_parts)

    def disconnect(self):
        """Close the persistent connection."""
        self.should_stop = True
        self._cleanup()

    def send_audio(self, audio_bytes: bytes):
        """Send raw PCM audio to Deepgram."""
        if not self.is_connected or not self.is_recording:
            return

        try:
            self.socket.send_media(audio_bytes)
        except Exception:
            self.is_connected = False

    def _cleanup(self):
        """Clean up connection resources."""
        if self.socket:
            try:
                self.socket._websocket.close()
            except Exception:
                pass

        if self.context_manager:
            try:
                self.context_manager.__exit__(None, None, None)
            except Exception:
                pass

        if self.listen_thread and self.listen_thread.is_alive():
            self.listen_thread.join(timeout=0.5)

        if self.keepalive_thread and self.keepalive_thread.is_alive():
            self.keepalive_thread.join(timeout=0.5)

        self.socket = None
        self.context_manager = None
        self.is_connected = False

    def _listen_loop(self):
        """Background thread that receives messages."""
        try:
            self.socket.start_listening()
        except Exception:
            pass
        finally:
            self.is_connected = False

    def _keepalive_loop(self):
        """Background thread that sends keep-alive messages when idle."""
        while not self.should_stop and self.is_connected:
            time.sleep(5)
            if self.is_connected and not self.is_recording:
                try:
                    self.socket._websocket.send(json.dumps({"type": "KeepAlive"}))
                except Exception:
                    break

    def _on_open(self, *args, **kwargs):
        self.is_connected = True

    def _on_message(self, *args, **kwargs):
        """Handle incoming transcription results."""
        with self.lock:
            accepting = self.is_recording or self._waiting_for_final
        if not accepting:
            return

        try:
            msg = args[0] if args else None
            if msg is None:
                return

            if hasattr(msg, "type") and msg.type == "Results":
                is_final = getattr(msg, "is_final", False)
                if hasattr(msg, "channel") and msg.channel.alternatives:
                    transcript = msg.channel.alternatives[0].transcript
                    if transcript and is_final:
                        with self.lock:
                            self.transcript_parts.append(transcript)
        except Exception:
            pass

    def _on_error(self, *args, **kwargs):
        error = args[0] if args else "Unknown error"
        if not self.should_stop:
            print(f"Deepgram error: {error}")

    def _on_close(self, *args, **kwargs):
        self.is_connected = False


def transcribe_with_whisper(audio_data: bytes) -> str:
    """Transcribe audio using OpenAI Whisper API."""
    if not audio_data:
        return ""

    client = OpenAI()
    audio_file = io.BytesIO(audio_data)
    audio_file.name = "audio.wav"

    transcription = client.audio.transcriptions.create(
        model="whisper-1",
        file=audio_file,
    )
    return transcription.text


def format_with_haiku(text: str) -> str:
    """Format transcribed text using Claude Haiku for code-aware cleanup."""
    try:
        client = anthropic.Anthropic()
        response = client.messages.create(
            model="claude-haiku-4-5-20251001",
            max_tokens=1024,
            system=HAIKU_SYSTEM,
            messages=[{"role": "user", "content": f"Format this transcription:\n\n{text}"}],
        )
        return response.content[0].text
    except Exception:
        return text


def paste_text(text: str):
    """Paste text at cursor position, preserving original clipboard contents."""
    if not text.strip():
        print("(empty transcription)")
        return

    # Save original clipboard
    try:
        original = pyperclip.paste()
    except Exception:
        original = ""

    # Copy transcribed text with trailing space and paste
    pyperclip.copy(text.strip() + " ")
    keyboard = Controller()
    keyboard.press(Key.cmd)
    keyboard.press("v")
    keyboard.release("v")
    keyboard.release(Key.cmd)

    # Wait for paste to complete, then restore original clipboard
    time.sleep(0.15)
    pyperclip.copy(original)

    preview = text[:50] + "..." if len(text) > 50 else text
    print(f"Pasted: {preview}")


class HotkeyHandler:
    """Handles hotkey detection and triggers recording/transcription."""

    def __init__(self, recorder: Recorder, transcriber=None):
        self.recorder = recorder
        self.transcriber = transcriber
        self.pressed_keys = set()
        self.is_recording = False
        self.use_haiku = False
        self.lock = threading.Lock()

    def on_press(self, key):
        with self.lock:
            if key not in HOTKEY_FORMAT:
                return

            self.pressed_keys.add(key)

            # Start recording when base combo is held
            if HOTKEY_BASE <= self.pressed_keys and not self.is_recording:
                self.use_haiku = Key.cmd in self.pressed_keys
                self._start_recording()

    def on_release(self, key):
        with self.lock:
            if key not in HOTKEY_FORMAT:
                return

            self.pressed_keys.discard(key)

            # Stop when any key from base combo is released
            if self.is_recording and key in HOTKEY_BASE:
                self._stop_recording()

    def _start_recording(self):
        """Start recording with appropriate backend."""
        transcriber = None

        if self.transcriber:
            # Auto-reconnect if connection dropped
            if not self.transcriber.is_connected:
                self.transcriber.connect()

            if self.transcriber.is_connected:
                self.transcriber.start_recording()
                transcriber = self.transcriber
            else:
                print("Deepgram reconnect failed, falling back to Whisper")

        if self.recorder.start(transcriber=transcriber):
            self.is_recording = True
            mode = "+Haiku" if self.use_haiku else "raw"
            backend = "Deepgram" if transcriber else "Whisper"
            print(f"Recording ({mode}, {backend})...")

    def _stop_recording(self):
        """Stop recording and process transcription in background."""
        self.is_recording = False
        use_haiku = self.use_haiku
        use_deepgram = self.transcriber and self.transcriber.is_recording
        audio_data = self.recorder.stop()

        threading.Thread(
            target=self._process_transcription,
            args=(audio_data, use_haiku, use_deepgram),
            daemon=True,
        ).start()

    def _process_transcription(self, audio_data: bytes, use_haiku: bool, use_deepgram: bool):
        """Process audio and paste transcribed text."""
        if use_deepgram and self.transcriber:
            text = self.transcriber.stop_recording()
        else:
            text = transcribe_with_whisper(audio_data)

        if use_haiku and text.strip():
            raw_text = text
            text = format_with_haiku(text)
            print(f"[raw] {raw_text}")
            print(f"[fmt] {text}")

        paste_text(text)


def setup():
    """Interactive setup to configure API keys."""
    print("Jotbot Setup")
    print("=" * 40)
    print(f"\nConfig will be saved to: {CONFIG_FILE}\n")

    # Create config directory
    CONFIG_DIR.mkdir(parents=True, exist_ok=True)

    # Load existing config if present
    existing = {}
    if CONFIG_FILE.exists():
        print("(Existing config found - press Enter to keep current values)\n")
        with open(CONFIG_FILE) as f:
            for line in f:
                line = line.strip()
                if line and not line.startswith("#") and "=" in line:
                    key, value = line.split("=", 1)
                    existing[key.strip()] = value.strip()

    config_lines = []

    # Transcription backend choice
    print("Choose transcription backend:")
    print("  1. Deepgram (recommended - real-time streaming, ~$0.0077/min)")
    print("  2. OpenAI Whisper (batch processing, ~$0.006/min)")

    current_backend = "1" if existing.get("USE_DEEPGRAM", "true").lower() == "true" else "2"
    choice = input(f"\nEnter 1 or 2 [{current_backend}]: ").strip() or current_backend
    use_deepgram = choice == "1"
    config_lines.append(f"USE_DEEPGRAM={'true' if use_deepgram else 'false'}")

    # API keys
    if use_deepgram:
        current = existing.get("DEEPGRAM_API_KEY", "")
        masked = f"{current[:8]}..." if len(current) > 8 else "(not set)"
        key = input(f"\nDeepgram API key [{masked}]: ").strip()
        config_lines.append(f"DEEPGRAM_API_KEY={key if key else current}")
    else:
        current = existing.get("OPENAI_API_KEY", "")
        masked = f"{current[:8]}..." if len(current) > 8 else "(not set)"
        key = input(f"\nOpenAI API key [{masked}]: ").strip()
        config_lines.append(f"OPENAI_API_KEY={key if key else current}")

    # Optional Anthropic key for Haiku formatting
    print("\n(Optional) Anthropic API key for Haiku formatting mode:")
    print("  This enables Cmd+Ctrl+Option to clean up transcriptions")
    current = existing.get("ANTHROPIC_API_KEY", "")
    masked = f"{current[:8]}..." if len(current) > 8 else "(not set)"
    key = input(f"Anthropic API key [{masked}]: ").strip()
    if key or current:
        config_lines.append(f"ANTHROPIC_API_KEY={key if key else current}")

    # Write config
    with open(CONFIG_FILE, "w") as f:
        f.write("\n".join(config_lines) + "\n")

    print(f"\nConfig saved to {CONFIG_FILE}")
    print("Run 'jotbot' to start dictation.")


def run():
    """Main dictation loop."""
    # Reload config in case setup was just run
    load_config()

    # Re-check USE_DEEPGRAM after loading config
    use_deepgram = os.environ.get("USE_DEEPGRAM", "").lower() in ("1", "true")

    # Validate required API keys
    if use_deepgram:
        if not os.environ.get("DEEPGRAM_API_KEY"):
            print("Error: DEEPGRAM_API_KEY not set.")
            print(f"Run 'jotbot --setup' to configure, or add it to {CONFIG_FILE}")
            sys.exit(1)
    elif not os.environ.get("OPENAI_API_KEY"):
        print("Error: OPENAI_API_KEY not set.")
        print(f"Run 'jotbot --setup' to configure, or add it to {CONFIG_FILE}")
        sys.exit(1)

    # Set up persistent Deepgram connection if enabled
    transcriber = None
    if use_deepgram:
        print("Connecting to Deepgram...")
        transcriber = DeepgramTranscriber()
        if transcriber.connect():
            print("Deepgram connected (persistent connection)")
        else:
            print("Warning: Deepgram connection failed, will use Whisper fallback")
            transcriber = None

    backend = "Deepgram streaming" if transcriber else "Whisper batch"
    print(f"Jotbot ready. (Backend: {backend})")
    print("  Ctrl+Option: dictate (raw)")
    print("  Cmd+Ctrl+Option: dictate with Haiku formatting (hold Cmd first)")

    recorder = Recorder()
    handler = HotkeyHandler(recorder, transcriber=transcriber)
    listener = Listener(on_press=handler.on_press, on_release=handler.on_release)
    listener.start()

    try:
        listener.join()
    except KeyboardInterrupt:
        if transcriber:
            transcriber.disconnect()


def main():
    parser = argparse.ArgumentParser(
        description="Push-to-talk voice dictation for macOS"
    )
    parser.add_argument(
        "--setup",
        action="store_true",
        help="Interactive setup to configure API keys"
    )
    args = parser.parse_args()

    if args.setup:
        setup()
    else:
        run()


if __name__ == "__main__":
    main()
