#!/usr/bin/env python3
"""
Stage 2 - Cozmo Chat Buddy
- Local Open WebUI / Ollama
- Faster-Whisper (local STT) with Google fallback
- eSpeak voice through Cozmo's speaker
- Cute head tilts while speaking
"""

import time
import threading
import subprocess
import tempfile
import os
import queue
from openai import OpenAI
import pycozmo
import speech_recognition as sr

# Optional local STT
try:
    from faster_whisper import WhisperModel
    LOCAL_STT_AVAILABLE = True
except ImportError:
    LOCAL_STT_AVAILABLE = False
    print("faster-whisper not installed – will use Google STT")

# ====================== CONFIG ======================
OPENWEBUI_BASE = "http://localhost:3000/api"          # or your Cloudflare URL
API_KEY = "MyAPIkey n sheeeeeeiit"
MODEL = "qwen3:8b or whatever"

SYSTEM_PROMPT = """You are Cozmo aka Little Buddy, a tiny adorable robot sitting on your charger.
You are curious, slightly chaotic, friendly, and a little dramatic. Keep answers short (1 sentence max). Speak like a cute little robot. and do not use * in your speech. You are chatting on hobostreamer.com
Never break character. Never mention you are an AI.
"""

# Voice settings (eSpeak)
ESPEAK_VOICE = "en+f3"        # higher / cuter
ESPEAK_SPEED = "155"

# Head movement
HEAD_BASE = 0.28
HEAD_TILT_AMOUNT = 0.16

# Listening
USE_LOCAL_STT = True          # set to False to force Google
WHISPER_MODEL_SIZE = "base"   # tiny / base / small
# ====================================================

client = OpenAI(base_url=OPENWEBUI_BASE, api_key=API_KEY)
recognizer = sr.Recognizer()
microphone = sr.Microphone()

# Load local Whisper if available
whisper_model = None
if LOCAL_STT_AVAILABLE and USE_LOCAL_STT:
    print("Loading local Whisper model...")
    whisper_model = WhisperModel(WHISPER_MODEL_SIZE, device="cpu", compute_type="int8")
    print("Local STT ready.")


def speak_espeak(text: str, cli):
    """Speak through Cozmo using eSpeak"""
    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
        wav_path = f.name

    try:
        subprocess.run([
            "espeak",
            "-v", ESPEAK_VOICE,
            "-s", ESPEAK_SPEED,
            "-w", wav_path,
            text
        ], check=True, capture_output=True)

        with open(wav_path, "rb") as audio_file:
            cli.play_audio(audio_file)

    except Exception as e:
        print(f"Cozmo audio error: {e}")
        # Fallback – play on PC so you can still hear it
        os.system(f'start /min "" "{wav_path}"')

    finally:
        try:
            os.unlink(wav_path)
        except:
            pass


def cute_head_tilt(cli, duration=3.0):
    """Gentle rhythmic head tilt while speaking"""
    start = time.time()
    direction = 1
    while time.time() - start < duration:
        angle = HEAD_BASE + (HEAD_TILT_AMOUNT * direction)
        try:
            cli.set_head_angle(angle)
        except:
            break
        time.sleep(0.30)
        direction *= -1
    try:
        cli.set_head_angle(HEAD_BASE)
    except:
        pass


def listen_once():
    """Listen and return transcribed text"""
    with microphone as source:
        print("\nListening... (speak now)")
        recognizer.adjust_for_ambient_noise(source, duration=0.35)
        try:
            audio = recognizer.listen(source, timeout=7, phrase_time_limit=10)
        except sr.WaitTimeoutError:
            print("No speech detected.")
            return None

    # Try local Whisper first
    if whisper_model is not None:
        try:
            with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
                wav_path = f.name
                with open(wav_path, "wb") as out:
                    out.write(audio.get_wav_data())

            segments, _ = whisper_model.transcribe(wav_path, language="en")
            text = " ".join([seg.text for seg in segments]).strip()
            os.unlink(wav_path)

            if text:
                print(f"You said (local): {text}")
                return text
        except Exception as e:
            print(f"Local STT failed: {e}")

    # Fallback to Google
    try:
        text = recognizer.recognize_google(audio)
        print(f"You said (Google): {text}")
        return text
    except sr.UnknownValueError:
        print("Didn't catch that.")
        return None
    except Exception as e:
        print(f"Recognition error: {e}")
        return None


def ask_llm(user_text: str) -> str:
    response = client.chat.completions.create(
        model=MODEL,
        messages=[
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": user_text}
        ],
        temperature=0.85,
        max_tokens=120
    )
    return response.choices[0].message.content.strip()


def main():
    print("Connecting to Cozmo...")
    with pycozmo.connect(enable_procedural_face=True) as cli:
        cli.enable_procedural_face(True)
        cli.set_head_angle(HEAD_BASE)
        time.sleep(0.5)

        print("Cozmo is ready on the charger.")
        print("Talk to him! (Ctrl+C to quit)\n")

        while True:
            try:
                user_text = listen_once()
                if not user_text:
                    continue

                print("Thinking...")
                reply = ask_llm(user_text)
                print(f"Cozmo: {reply}\n")

                # Head tilt + speak at the same time
                tilt_thread = threading.Thread(
                    target=cute_head_tilt,
                    args=(cli, 3.2),
                    daemon=True
                )
                tilt_thread.start()

                speak_espeak(reply, cli)
                tilt_thread.join(timeout=4.0)

                time.sleep(0.3)

            except KeyboardInterrupt:
                print("\nBye little buddy!")
                break
            except Exception as e:
                print(f"Error: {e}")
                time.sleep(1)


if __name__ == "__main__":
    main()