import asyncio
import os
import tempfile
from pathlib import Path

import edge_tts
import gradio as gr
from faster_whisper import WhisperModel
from huggingface_hub import hf_hub_download
from llama_cpp import Llama


MODEL_REPO = os.getenv(
    "ADI_MODEL_REPO",
    "AdvancedDataIntelligence/adi-qwen3.5-4b-glm5.2-general-GGUF",
)
MODEL_FILE = os.getenv(
    "ADI_MODEL_FILE",
    "adi-qwen3.5-4b-glm5.2-general-q4_k_m.gguf",
)
WHISPER_MODEL = os.getenv("ADI_WHISPER_MODEL", "tiny.en")
TTS_VOICE = os.getenv("ADI_TTS_VOICE", "en-US-AriaNeural")

SYSTEM_PROMPT = (
    "You are ADI (Advanced Data Intelligence), a concise voice assistant. "
    "Reply naturally in short spoken answers unless the user asks for detail."
)

_llm = None
_stt = None


def get_llm():
    global _llm
    if _llm is None:
        model_path = hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILE)
        _llm = Llama(
            model_path=model_path,
            n_ctx=4096,
            n_threads=max(2, min(4, os.cpu_count() or 2)),
            chat_format="chatml",
            verbose=False,
        )
    return _llm


def get_stt():
    global _stt
    if _stt is None:
        _stt = WhisperModel(WHISPER_MODEL, device="cpu", compute_type="int8")
    return _stt


def transcribe(audio_path):
    if not audio_path:
        return ""

    segments, _info = get_stt().transcribe(
        audio_path,
        beam_size=1,
        vad_filter=True,
    )
    return " ".join(segment.text.strip() for segment in segments).strip()


def chat_once(message, history, temperature, max_tokens):
    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
    messages.extend(history or [])
    messages.append({"role": "user", "content": message})

    stream = get_llm().create_chat_completion(
        messages=messages,
        temperature=float(temperature),
        max_tokens=int(max_tokens),
        stream=True,
    )

    response = ""
    for chunk in stream:
        delta = chunk["choices"][0]["delta"].get("content", "")
        if delta:
            response += delta
    return response.strip()


async def speak_async(text):
    output_path = Path(tempfile.NamedTemporaryFile(suffix=".mp3", delete=False).name)
    communicate = edge_tts.Communicate(text, TTS_VOICE)
    await communicate.save(str(output_path))
    return str(output_path)


def speak(text):
    if not text.strip():
        return None
    return asyncio.run(speak_async(text))


def respond(audio_path, typed_message, history, temperature, max_tokens):
    history = history or []
    typed_message = (typed_message or "").strip()
    transcript = typed_message or transcribe(audio_path)

    if not transcript:
        return "", history, None, "Record audio or type a message first."

    reply = chat_once(transcript, history, temperature, max_tokens)
    next_history = history + [
        {"role": "user", "content": transcript},
        {"role": "assistant", "content": reply},
    ]
    audio_reply = speak(reply)
    return transcript, next_history, audio_reply, "Ready"


def clear_chat():
    return "", [], None, "Ready"


with gr.Blocks(
    title="ADI Voice Demo",
    fill_height=True,
    css="""
    .gradio-container { max-width: 1120px !important; margin: auto !important; }
    #status-box textarea { font-size: 0.9rem; }
    """,
) as demo:
    gr.Markdown("# ADI Voice Demo")

    history_state = gr.State([])

    with gr.Row():
        with gr.Column(scale=1, min_width=320):
            mic = gr.Audio(
                sources=["microphone", "upload"],
                type="filepath",
                label="Speak to ADI",
            )
            typed = gr.Textbox(
                label="Or type",
                placeholder="Say hello, ask a question, or paste text here.",
                lines=3,
            )
            with gr.Row():
                submit = gr.Button("Talk", variant="primary")
                clear = gr.Button("Clear")
            temperature = gr.Slider(
                0.0,
                1.5,
                value=0.7,
                step=0.1,
                label="Temperature",
            )
            max_tokens = gr.Slider(
                32,
                512,
                value=160,
                step=32,
                label="Max tokens",
            )

        with gr.Column(scale=2, min_width=420):
            chatbot = gr.Chatbot(
                label="Conversation",
                height=460,
                autoscroll=True,
            )
            transcript = gr.Textbox(label="Transcript", interactive=False)
            voice = gr.Audio(label="ADI voice", autoplay=True, type="filepath")
            status = gr.Textbox(
                label="Status",
                value="Ready",
                interactive=False,
                elem_id="status-box",
            )

    submit.click(
        respond,
        inputs=[mic, typed, history_state, temperature, max_tokens],
        outputs=[transcript, history_state, voice, status],
    ).then(
        lambda h: h,
        inputs=history_state,
        outputs=chatbot,
    )

    clear.click(
        clear_chat,
        outputs=[transcript, history_state, voice, status],
    ).then(
        lambda h: h,
        inputs=history_state,
        outputs=chatbot,
    )


if __name__ == "__main__":
    demo.launch()