import re

import ollama

from audio.tts import speak

AI_MODEL = "gemma4:e2b-it-qat"
AI_HOST = "http://localhost:11434"
def clean_response(text):
    cleaned = re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL | re.IGNORECASE)
    lines = [line for line in cleaned.splitlines() if line.strip()]
    return "\n".join(lines).strip()
def ask_ai(prompt):
    try:
        full_prompt = (
            "You must answer in exactly one or two sentences.\n\n"
            f"Question: {prompt}\n"
            "Answer:"
        )
        print("Thinking...", end="", flush=True)
        client = ollama.Client(host=AI_HOST)
        stream = client.generate(
            model=AI_MODEL,
            prompt=full_prompt,
            stream=True
        )
        ai_text = ""
        for chunk in stream:
            text_piece = chunk.get("response", "")
            if text_piece:
                print(text_piece, end="", flush=True)
                ai_text += text_piece
        print()
        final_text = clean_response(ai_text)
        if not final_text:
            speak("Sorry, I got no response from the AI.")
        else:
            speak(final_text)
    except Exception as e:

        print("\nAI Error:", e)

        speak("Sorry, I couldn't connect to the AI.")