import argparse
from io import BytesIO
from itertools import chain
from pathlib import Path
import sqlite3
from time import perf_counter
import wave

from faster_whisper import WhisperModel
from piper import PiperVoice
from ask_handbook import select_quote
from handbook_search import DEFAULT_DB, search
from transcribe_notes import transcribe


def write_reply(voice, text, path, start):
    chunks = iter(voice.synthesize(text))
    first = next(chunks, None)
    if first is None:
        raise ValueError("Piper returned no audio.")
    first_chunk = perf_counter() - start
    with BytesIO() as buffer:
        with wave.open(buffer, "wb") as audio:
            audio.setnchannels(first.sample_channels)
            audio.setsampwidth(first.sample_width)
            audio.setframerate(first.sample_rate)
            for chunk in chain((first,), chunks):
                audio.writeframesraw(chunk.audio_int16_bytes)
        with path.open("xb") as output:
            output.write(buffer.getbuffer())
    return first_chunk


def run(args):
    setup = perf_counter()
    voice = PiperVoice.load(args.voice)
    recognizer = (WhisperModel("tiny.en", device="cpu", compute_type="int8", local_files_only=True)
                  if args.audio else None)
    start = perf_counter()
    print(f"Load speech models: {start - setup:.2f}s")

    if recognizer:
        lines, _ = transcribe(recognizer, args.audio, vad_filter=True)
        question = " ".join(text for _, text in lines)
        print(f"Transcription: {perf_counter() - start:.2f}s")
    else:
        question = args.question.strip()
    if not question:
        print("No question detected. Try a clear recording or pass --question.")
        return
    print(f"Question: {question}")

    retrieval_start = perf_counter()
    passages = search(args.db, question)
    generation_start = perf_counter()
    print(f"Retrieval: {generation_start - retrieval_start:.2f}s")
    print("Candidates:", ", ".join(f"{item.source} p.{item.page}" for item in passages) or "none")
    result = select_quote(passages, question, args.model)
    print(f"Quote selection and check: {perf_counter() - generation_start:.2f}s")
    if result is None:
        spoken = "I could not verify a quote in the handbook."
    else:
        passage, spoken = result
        print(f"Source: {passage.source}, page {passage.page}")
    print(f"Spoken text: {spoken}")
    first_chunk = write_reply(voice, spoken, args.output, start)
    print(f"First audio chunk ready: {first_chunk:.2f}s")
    print(f"WAV ready: {perf_counter() - start:.2f}s | {args.output}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Ask a local handbook question and save a spoken reply.")
    inputs = parser.add_mutually_exclusive_group(required=True)
    inputs.add_argument("--audio", type=Path)
    inputs.add_argument("--question")
    parser.add_argument("--voice", type=Path, default=Path("voices/en_US-lessac-medium.onnx"))
    parser.add_argument("--db", type=Path, default=DEFAULT_DB)
    parser.add_argument("--model", default="qwen2.5:1.5b")
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()
    if args.output.exists():
        parser.error("Choose a new output file.")
    if args.audio and not args.audio.is_file():
        parser.error(f"Recording not found: {args.audio}")
    try:
        run(args)
    except (OSError, KeyError, RuntimeError, sqlite3.Error, ValueError) as error:
        raise SystemExit(str(error)) from None