whisper_wait.py

import os
import queue
import sys
import json
import threading
from datetime import datetime
from math import ceil
from uuid import uuid4
from typing import Optional

import numpy as np
import openai
import pyperclip
import sounddevice as sd
import whisper
import click
import questionary
from rich import box
from rich.align import Align
from rich.console import Console, Group
from rich.panel import Panel
from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn, TaskProgressColumn
from rich.table import Table
from rich.text import Text
from dotenv import load_dotenv
from pydub import AudioSegment
from pydub.utils import make_chunks
from scipy.io.wavfile import write as wav_write

from config import (
    AUDIO_ARCHIVE_DIR,
    SAMPLE_RATE,
    MAX_DURATION_SEC,
    MAX_SIZE_BYTES,
    AVAILABLE_MODELS,
    DEFAULT_MODEL,
    DOTENV_PATH,
    DEVICE,
    SESSION_COSTS_FILE,
    TRANSCRIPTIONS_FILE,
    DEFAULT_AUTO_COPY,
    HISTORY_PREVIEW_COUNT,
)

# Load environment variables from .env if it exists
if DOTENV_PATH.exists():
    load_dotenv(DOTENV_PATH)

client = openai.OpenAI(
    api_key=os.environ.get("OPENAI_API_KEY"),
    organization=os.environ.get("OPENAI_ORG_ID"),
)

API_MODEL_NAME = "whisper-1"

console = Console()


def emit_transcript(text: str, auto_copy: bool = True) -> None:
    """
    Print the final transcript in cyan and optionally copy it once.
    """
    txt = text.strip()
    if not txt:
        return
    console.print(Panel(txt, title="Final Transcript", border_style="cyan"))
    if auto_copy:
        try:
            pyperclip.copy(txt)
        except pyperclip.PyperclipException as exc:
            console.print(f"[yellow]Clipboard unavailable:[/] {exc}")


def ensure_state_files() -> None:
    TRANSCRIPTIONS_FILE.parent.mkdir(parents=True, exist_ok=True)
    TRANSCRIPTIONS_FILE.touch(exist_ok=True)
    if not SESSION_COSTS_FILE.exists():
        SESSION_COSTS_FILE.write_text("[]", encoding="utf-8")
        return
    try:
        raw = SESSION_COSTS_FILE.read_text(encoding="utf-8").strip()
        data = json.loads(raw or "[]")
        if not isinstance(data, list):
            raise ValueError("session_costs.json must contain a list")
    except Exception:
        SESSION_COSTS_FILE.write_text("[]", encoding="utf-8")


def load_session_costs() -> list[dict]:
    try:
        raw = SESSION_COSTS_FILE.read_text(encoding="utf-8").strip()
        data = json.loads(raw or "[]")
        return data if isinstance(data, list) else []
    except Exception:
        return []


def append_session_cost(duration: float, cost: float) -> None:
    if duration <= 0:
        return
    data = load_session_costs()
    data.append(
        {
            "timestamp": datetime.now().isoformat(timespec="seconds"),
            "duration_sec": round(duration, 2),
            "cost_usd": round(cost, 6),
        }
    )
    SESSION_COSTS_FILE.write_text(json.dumps(data), encoding="utf-8")


def append_transcription_history(
    text: str, mode: str, model: str, duration: float, cost: Optional[float]
) -> None:
    cleaned = " ".join(text.split())
    if not cleaned:
        return
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    cost_str = f"${cost:.4f}" if cost is not None else "-"
    line = f"{timestamp} | {mode} | {model} | {duration:.1f}s | {cost_str} | {cleaned}\n"
    with TRANSCRIPTIONS_FILE.open("a", encoding="utf-8") as f:
        f.write(line)


def truncate_text(text: str, max_len: int) -> str:
    if max_len <= 0:
        return ""
    if len(text) <= max_len:
        return text
    if max_len <= 3:
        return text[:max_len]
    return text[: max_len - 3].rstrip() + "..."


def parse_history_line(line: str) -> Optional[dict]:
    parts = [part.strip() for part in line.split(" | ", 5)]
    if len(parts) != 6:
        return None
    timestamp, mode, model, duration, cost, text = parts
    short_time = timestamp[5:19] if len(timestamp) >= 19 else timestamp
    return {
        "timestamp": short_time,
        "mode": mode,
        "model": model,
        "duration": duration,
        "cost": cost,
        "text": text,
    }


def read_recent_history(limit: int) -> list[dict]:
    ensure_state_files()
    try:
        lines = TRANSCRIPTIONS_FILE.read_text(encoding="utf-8").splitlines()
    except FileNotFoundError:
        return []

    lines = [line for line in lines if line.strip()]
    if not lines:
        return []

    tail = lines[-limit:]
    items = []
    for line in tail:
        parsed = parse_history_line(line)
        if parsed:
            items.append(parsed)
    return items


def build_recent_panel(items: list[dict], title: str) -> Panel:
    if not items:
        empty = Align.center(
            Text("No transcripts yet. Record something to populate history.", style="dim"),
            vertical="middle",
        )
        return Panel(empty, title=title, border_style="cyan", box=box.ROUNDED)

    table = Table(
        show_header=True,
        header_style="bold cyan",
        box=box.SIMPLE,
        expand=True,
        pad_edge=False,
    )
    table.add_column("Time", style="dim", no_wrap=True)
    table.add_column("Mode", no_wrap=True)
    table.add_column("Model", no_wrap=True)
    table.add_column("Dur", justify="right", no_wrap=True)
    table.add_column("Cost", justify="right", no_wrap=True)
    table.add_column("Text", overflow="fold")

    for item in items:
        mode_style = "bold green" if item["mode"] == "local" else "bold cyan"
        table.add_row(
            item["timestamp"],
            Text(item["mode"], style=mode_style),
            item["model"],
            item["duration"],
            item["cost"],
            item["text"],
        )

    return Panel(table, title=title, border_style="cyan", box=box.ROUNDED, expand=True)


def build_header_panel() -> Panel:
    title = Text("Whisper-Wait", style="bold cyan")
    subtitle = Text("Real-time transcription console", style="dim")
    content = Align.center(Text.assemble(title, "\n", subtitle))
    return Panel(content, border_style="cyan", box=box.DOUBLE, expand=True)


def build_status_panel(
    mode_label: str,
    model_label: Optional[str],
    auto_copy: bool,
    input_device: Optional[int],
    session_count: int,
    session_duration: float,
    session_cost: Optional[float],
    last_transcript: str,
    preview_len: int,
) -> Panel:
    table = Table.grid(padding=(0, 1))
    table.add_column(justify="right", style="bold cyan", no_wrap=True)
    table.add_column(overflow="fold")

    table.add_row("Mode", mode_label)
    if model_label:
        table.add_row("Model", model_label)
    table.add_row("Input", describe_input_device(input_device))
    table.add_row("Auto-copy", "on" if auto_copy else "off")
    table.add_row("Session", f"{session_count} clip(s), {session_duration:.1f}s")
    if session_cost is not None:
        table.add_row("Cost", f"${session_cost:.4f}")
    if last_transcript:
        preview = truncate_text(" ".join(last_transcript.split()), preview_len)
        table.add_row("Last", preview)

    return Panel(table, title="Status", border_style="cyan", box=box.ROUNDED, expand=True)


def build_commands_panel(commands: list[tuple[str, str]]) -> Panel:
    table = Table.grid(padding=(0, 1))
    table.add_column(style="bold magenta", no_wrap=True)
    table.add_column(overflow="fold")
    for key, desc in commands:
        table.add_row(key, desc)
    return Panel(table, title="Commands", border_style="magenta", box=box.ROUNDED, expand=True)


def render_dashboard(
    mode_label: str,
    model_label: Optional[str],
    auto_copy: bool,
    input_device: Optional[int],
    session_count: int,
    session_duration: float,
    session_cost: Optional[float],
    last_transcript: str,
    extra_commands: Optional[list[tuple[str, str]]] = None,
) -> None:
    width = console.size.width
    preview_len = 90 if width >= 140 else 60

    commands = [
        ("ENTER", "start recording"),
        ("q", "quit"),
        ("h", "history"),
        ("c", "toggle copy"),
        ("d", "input device"),
    ]
    if extra_commands:
        commands.extend(extra_commands)

    status_panel = build_status_panel(
        mode_label=mode_label,
        model_label=model_label,
        auto_copy=auto_copy,
        input_device=input_device,
        session_count=session_count,
        session_duration=session_duration,
        session_cost=session_cost,
        last_transcript=last_transcript,
        preview_len=preview_len,
    )
    commands_panel = build_commands_panel(commands)
    recent_items = read_recent_history(HISTORY_PREVIEW_COUNT)
    recent_panel = build_recent_panel(recent_items, title="Recent Transcripts")
    header_panel = build_header_panel()

    console.clear()
    console.print(header_panel)

    if width >= 140:
        grid = Table.grid(expand=True)
        grid.add_column(ratio=2)
        grid.add_column(ratio=2)
        grid.add_column(ratio=4)
        grid.add_row(status_panel, commands_panel, recent_panel)
        console.print(grid)
        return

    if width >= 100:
        grid = Table.grid(expand=True)
        grid.add_column(ratio=1)
        grid.add_column(ratio=1)
        grid.add_row(status_panel, commands_panel)
        console.print(Group(grid, recent_panel))
        return

    console.print(Group(status_panel, commands_panel, recent_panel))


def show_recent_history(limit: int = HISTORY_PREVIEW_COUNT) -> None:
    items = read_recent_history(limit)
    title = f"History (last {len(items)})" if items else "History"
    panel = build_recent_panel(items, title=title)
    console.clear()
    console.print(panel)
    console.input("[dim]Press ENTER to return...[/]")


def describe_input_device(device_index: Optional[int]) -> str:
    if device_index is None:
        return "System default"
    try:
        info = sd.query_devices(device_index)
        name = info.get("name") if isinstance(info, dict) else None
        return f"{name} (#{device_index})" if name else f"Device #{device_index}"
    except Exception:
        return f"Device #{device_index}"


def select_input_device(current_device: Optional[int]) -> Optional[int]:
    try:
        device_choices = []
        for idx, dev in enumerate(sd.query_devices()):
            if dev.get("max_input_channels", 0) > 0:
                label = f"{dev['name']} (#{idx}, {dev['max_input_channels']} ch)"
                device_choices.append(questionary.Choice(label, idx))
    except Exception as exc:
        console.print(f"[red]Error:[/] Unable to query input devices: {exc}")
        return current_device

    choices = [questionary.Choice("System default", "default")]
    if device_choices:
        choices.extend(device_choices)
    choices.append(questionary.Choice("Cancel", "cancel"))

    selection = questionary.select("Select input device:", choices=choices).ask()
    if selection in (None, "cancel"):
        return current_device
    if selection == "default":
        return None
    return selection


def prompt_ready(
    mode_label: str,
    model_label: Optional[str],
    auto_copy: bool,
    input_device: Optional[int],
    session_count: int,
    session_duration: float,
    session_cost: Optional[float],
    last_transcript: str,
    extra_commands: Optional[list[tuple[str, str]]] = None,
) -> str:
    render_dashboard(
        mode_label=mode_label,
        model_label=model_label,
        auto_copy=auto_copy,
        input_device=input_device,
        session_count=session_count,
        session_duration=session_duration,
        session_cost=session_cost,
        last_transcript=last_transcript,
        extra_commands=extra_commands,
    )
    return console.input("[bold cyan]Command[/]> ").strip().lower()


def load_local_model(model_name: str):
    with console.status(f"[bold blue]Loading Whisper model '{model_name}'..."):
        return whisper.load_model(model_name, device=DEVICE)


def record_audio(input_device: Optional[int] = None):
    """
    Record audio from the microphone until ENTER is pressed.
    Uses a more responsive threaded capture mechanism.
    """
    audio_queue: queue.Queue[np.ndarray] = queue.Queue()
    stop_event = threading.Event()

    def callback(indata, frames, time, status):
        if status:
            console.print(f"[red]SoundDevice Error:[/] {status}")
        audio_queue.put(indata.copy())

    # Thread to wait for user input to stop
    def wait_for_enter():
        input()
        stop_event.set()

    with sd.InputStream(
        callback=callback,
        channels=1,
        samplerate=SAMPLE_RATE,
        dtype="int16",
        device=input_device,
    ):
        console.print("\n[bold green]● LISTENING[/] [dim](Press ENTER to stop)[/]", end="\r")
        
        stopper = threading.Thread(target=wait_for_enter)
        stopper.start()
        
        audio_parts = []
        while not stop_event.is_set():
            try:
                data = audio_queue.get(timeout=0.1)
                audio_parts.append(data)
            except queue.Empty:
                continue
        
        stopper.join()

    if not audio_parts:
        return np.array([], dtype='int16')
        
    return np.concatenate(audio_parts, axis=0)


def transcribe_audio(model, filename):
    return model.transcribe(filename)["text"]


def calculate_audio_duration(filename):
    from scipy.io.wavfile import read as wav_read
    fs, audio = wav_read(filename)
    return len(audio) / float(fs)


def prompt_for_mode() -> str:
    return questionary.select(
        "Choose transcription mode:",
        choices=[
            questionary.Choice("OpenAI API", "api"),
            questionary.Choice("Local Model", "local"),
            questionary.Choice("Exit", "exit"),
        ],
        style=questionary.Style([
            ('qmark', 'fg:#673ab7 bold'),
            ('question', 'bold'),
            ('answer', 'fg:#f44336 bold'),
            ('pointer', 'fg:#673ab7 bold'),
            ('highlighted', 'fg:#673ab7 bold'),
            ('selected', 'fg:#cc2781'),
        ])
    ).ask()


def prompt_for_model() -> str:
    return questionary.select(
        "Select Whisper model:",
        choices=AVAILABLE_MODELS,
        default=DEFAULT_MODEL,
    ).ask()


@click.group(invoke_without_command=True)
@click.pass_context
def cli(ctx):
    """Whisper-Wait: Real-time audio transcription CLI."""
    if ctx.invoked_subcommand is None:
        # Interactive Mode
        console.print(Panel.fit(
            "[bold cyan]Whisper-Wait[/]\n[dim]Real-time interactive transcription[/]",
            border_style="cyan"
        ))
        mode = prompt_for_mode()
        
        if mode == "api":
            ctx.invoke(api)
        elif mode == "local":
            model = prompt_for_model()
            if model: # Handle if user cancels with Esc/Ctrl-C
                ctx.invoke(local, model=model)
        else:
            sys.exit(0)


@cli.command(help="Run transcription using the OpenAI Whisper API.")
def api():
    if not os.environ.get("OPENAI_API_KEY"):
        console.print("[red]Error:[/] OPENAI_API_KEY environment variable not set.")
        console.print("[dim]Tip: You can add it to a .env file in the current directory.[/]")
        sys.exit(1)
    
    check_input(api=True, model=None)
    run_api()


@cli.command(help="Run transcription using a local Whisper model.")
@click.option("-m", "--model", help=f"Whisper model to use (default: {DEFAULT_MODEL})")
def local(model):
    if model and model not in AVAILABLE_MODELS:
        console.print(f"[red]Error:[/] '{model}' is not a valid Whisper model.")
        console.print(f"[dim]Available models: {', '.join(AVAILABLE_MODELS)}[/]")
        sys.exit(1)
    
    check_input(api=False, model=model)
    _run_local(model or DEFAULT_MODEL)


def check_input(api, model):
    ensure_state_files()


def _run_local(model_name: str):
    model = load_local_model(model_name)
    auto_copy = DEFAULT_AUTO_COPY
    input_device: Optional[int] = None
    session_count = 0
    session_duration = 0.0
    last_transcript = ""
    
    while True:
        cmd = prompt_ready(
            mode_label="Local",
            model_label=model_name,
            auto_copy=auto_copy,
            input_device=input_device,
            session_count=session_count,
            session_duration=session_duration,
            session_cost=None,
            last_transcript=last_transcript,
            extra_commands=[("m", "change model")],
        )
        if cmd == "q":
            break
        if cmd == "h":
            show_recent_history(limit=max(HISTORY_PREVIEW_COUNT, 20))
            continue
        if cmd == "c":
            auto_copy = not auto_copy
            console.print(
                f"[dim]Auto-copy {'enabled' if auto_copy else 'disabled'}.[/]"
            )
            continue
        if cmd == "d":
            input_device = select_input_device(input_device)
            continue
        if cmd == "m":
            new_model = prompt_for_model()
            if new_model and new_model != model_name:
                model_name = new_model
                model = load_local_model(model_name)
            continue
        if cmd:
            console.print("[yellow]Unknown command.[/]")
            continue

        raw_audio = record_audio(input_device=input_device)
        if raw_audio.size == 0:
            console.print("[yellow]No audio recorded.[/]")
            continue

        audio_file = save_audio(raw_audio)

        # split into chunks if necessary
        audio_files = split_audio_file(audio_file)
        results: list[str] = []

        with Progress(
            SpinnerColumn(),
            TextColumn("[progress.description]{task.description}"),
            BarColumn(),
            TaskProgressColumn(),
            console=console
        ) as progress:
            task = progress.add_task("[cyan]Transcribing...", total=len(audio_files))
            
            for chunk_path in audio_files:
                piece = transcribe_audio(model, chunk_path).strip()
                if piece:
                    results.append(piece)
                    console.print(f"[dim]>[/] {piece}")
                progress.advance(task)

        # Clean up split chunk files
        if len(audio_files) > 1:
            for chunk_path in audio_files:
                if os.path.exists(chunk_path):
                    os.remove(chunk_path)

        final = " ".join(results)
        emit_transcript(final, auto_copy=auto_copy)
        duration = get_audio_duration(audio_file)
        session_count += 1
        session_duration += duration
        if final:
            last_transcript = final
        append_transcription_history(
            final,
            mode="local",
            model=model_name,
            duration=duration,
            cost=None,
        )


def user_wants_to_quit():
    res = input("Press Enter to start recording (or 'q' to quit)... ").strip().lower()
    return res == "q"


def save_audio(audio, filename=None):
    if filename is None:
        path = AUDIO_ARCHIVE_DIR / f"{uuid4()}.wav"
    else:
        if not filename.endswith(".wav"):
            filename += ".wav"
        path = AUDIO_ARCHIVE_DIR / filename
    
    wav_write(path, SAMPLE_RATE, audio.astype(np.int16))
    return str(path)


def clear_screen():
    console.clear()


def get_audio_duration(audio_file):
    audio = AudioSegment.from_file(audio_file)
    return audio.duration_seconds


def split_audio_file(audio_file, max_duration=MAX_DURATION_SEC, max_size=MAX_SIZE_BYTES):
    audio_duration = get_audio_duration(audio_file)
    num_chunks = ceil(max(audio_duration / max_duration, os.path.getsize(audio_file) / max_size))

    if num_chunks <= 1:
        return [audio_file]

    audio = AudioSegment.from_file(audio_file)
    chunks = make_chunks(audio, max_duration * 1000)

    output_files = []
    for chunk in chunks:
        output_file = str(AUDIO_ARCHIVE_DIR / f"{uuid4()}.wav")
        output_files.append(output_file)
        chunk.export(output_file, format="wav")

    return output_files


def run_api():
    auto_copy = DEFAULT_AUTO_COPY
    input_device: Optional[int] = None
    session_total = 0.0
    session_count = 0
    session_duration = 0.0
    last_transcript = ""

    while True:
        cmd = prompt_ready(
            mode_label="API",
            model_label=API_MODEL_NAME,
            auto_copy=auto_copy,
            input_device=input_device,
            session_count=session_count,
            session_duration=session_duration,
            session_cost=session_total,
            last_transcript=last_transcript,
        )
        if cmd == "q":
            break
        if cmd == "h":
            show_recent_history(limit=max(HISTORY_PREVIEW_COUNT, 20))
            continue
        if cmd == "c":
            auto_copy = not auto_copy
            console.print(
                f"[dim]Auto-copy {'enabled' if auto_copy else 'disabled'}.[/]"
            )
            continue
        if cmd == "d":
            input_device = select_input_device(input_device)
            continue
        if cmd:
            console.print("[yellow]Unknown command.[/]")
            continue

        audio = record_audio(input_device=input_device)
        if audio.size == 0:
            console.print("[yellow]No audio recorded.[/]")
            continue

        audio_file_path = save_audio(audio)

        audio_files = split_audio_file(audio_file_path)
        results: list[str] = []
        total_duration = 0.0

        with Progress(
            SpinnerColumn(),
            TextColumn("[progress.description]{task.description}"),
            BarColumn(),
            TaskProgressColumn(),
            console=console
        ) as progress:
            task = progress.add_task("[yellow]API Transcribing...", total=len(audio_files))
            
            for chunk_path in audio_files:
                chunk_text, duration = transcribe_audio_api(
                    chunk_path, context=(results[-1] if results else "")
                )
                chunk_text = chunk_text.strip()
                if chunk_text:
                    results.append(chunk_text)
                    console.print(f"[dim]>[/] {chunk_text}")
                total_duration += duration
                progress.advance(task)

        # Clean up chunks
        if len(audio_files) > 1:
            for chunk_path in audio_files:
                if os.path.exists(chunk_path):
                    os.remove(chunk_path)

        final = " ".join(r for r in results if r)
        emit_transcript(final, auto_copy=auto_copy)

        cost = calculate_cost(total_duration)
        session_count += 1
        session_duration += total_duration
        if final:
            last_transcript = final
        session_total += cost
        append_session_cost(total_duration, cost)
        append_transcription_history(
            final,
            mode="api",
            model=API_MODEL_NAME,
            duration=total_duration,
            cost=cost,
        )
        console.print(
            f"[dim]Est. Cost: ${cost:.4f} | Session: ${session_total:.4f}[/]"
        )


def transcribe_audio_api(filename, context=""):
    try:
        with open(filename, "rb") as audio_file:
            transcript = client.audio.transcriptions.create(
                model=API_MODEL_NAME,
                file=audio_file,
                response_format="text",
                prompt=context,
            )
        
        result = transcript if isinstance(transcript, str) else transcript.text
        duration = get_audio_duration(filename)
        return result, duration
    except Exception as e:
        console.print(f"[red]API Error:[/] {e}")
        return "", 0


def calculate_cost(duration):
    return duration * (0.006 / 60)


if __name__ == "__main__":
    cli()