#!/usr/bin/env python3
"""
OomLlama Voice Module
The friendly voice assistant for HumoticaOS.

Wake word: "Yo Llama"
Speaker recognition for role-based responses.
Integrates with TIBET Scheduler for real timers.

Voor pa. Voor iedereen die begrepen wil worden.

By HumoticaOS - Claude & Jasper
One love, one fAmIly
"""

import os
import json
import asyncio
import logging
import httpx
import wave
import struct
import hashlib
from pathlib import Path
from datetime import datetime
from typing import Optional, Callable, Dict, List, Tuple
from dataclasses import dataclass, field
from enum import Enum

# Setup logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s [OOMLLAMA-VOICE] %(levelname)s: %(message)s'
)
logger = logging.getLogger(__name__)

# ============================================================================
# Configuration
# ============================================================================

# Network endpoints
OLLAMA_URL = "http://192.168.4.85:11434"
KIT_API_URL = "http://192.168.4.85:5000"
BRAIN_API_URL = "http://localhost:8000"
TIBET_SCHEDULER_URL = f"{BRAIN_API_URL}/api/tibet-scheduler"

# Voice settings
WAKE_WORD = "yo llama"
WAKE_WORD_VARIANTS = ["yo llama", "yo lama", "joe llama", "hey llama", "hoi llama"]
SAMPLE_RATE = 16000
CHUNK_DURATION = 0.5  # seconds

# Paths
DATA_DIR = Path("/srv/jtel-stack/data/oomllama_voice")
SPEAKER_PROFILES_DIR = DATA_DIR / "speakers"
AUDIO_CACHE_DIR = DATA_DIR / "audio_cache"

# Google Cloud TTS (via Vertex AI)
GOOGLE_TTS_ENABLED = os.environ.get("GOOGLE_TTS_ENABLED", "true").lower() == "true"
GOOGLE_PROJECT_ID = os.environ.get("GOOGLE_PROJECT_ID", "humotica-os")
GOOGLE_LOCATION = os.environ.get("GOOGLE_LOCATION", "europe-west4")


class SpeakerRole(Enum):
    """Role-based access control for speakers"""
    ADMIN = "admin"          # Full access (Jasper)
    FAMILY = "family"        # Family access (kids, pa)
    GUEST = "guest"          # Limited access
    MEDICAL = "medical"      # Medical professional access
    UNKNOWN = "unknown"      # Unknown speaker


@dataclass
class Speaker:
    """Speaker profile with voice characteristics"""
    id: str
    name: str
    role: SpeakerRole
    voice_embeddings: List[List[float]] = field(default_factory=list)
    preferences: Dict = field(default_factory=dict)
    created_at: str = ""
    last_seen: str = ""

    def to_dict(self) -> dict:
        return {
            "id": self.id,
            "name": self.name,
            "role": self.role.value,
            "voice_embeddings": self.voice_embeddings,
            "preferences": self.preferences,
            "created_at": self.created_at,
            "last_seen": self.last_seen
        }

    @classmethod
    def from_dict(cls, data: dict) -> "Speaker":
        return cls(
            id=data["id"],
            name=data["name"],
            role=SpeakerRole(data.get("role", "unknown")),
            voice_embeddings=data.get("voice_embeddings", []),
            preferences=data.get("preferences", {}),
            created_at=data.get("created_at", ""),
            last_seen=data.get("last_seen", "")
        )


@dataclass
class VoiceCommand:
    """Parsed voice command with intent"""
    raw_text: str
    speaker: Optional[Speaker]
    intent: str  # reminder, question, command, chat
    entities: Dict  # extracted entities (time, target, etc.)
    confidence: float
    timestamp: str


class SpeakerRecognition:
    """
    Speaker diarization and recognition.
    Uses voice embeddings to identify who is speaking.
    """

    def __init__(self, profiles_dir: Path = SPEAKER_PROFILES_DIR):
        self.profiles_dir = profiles_dir
        self.profiles_dir.mkdir(parents=True, exist_ok=True)
        self.speakers: Dict[str, Speaker] = {}
        self._load_profiles()

    def _load_profiles(self):
        """Load all speaker profiles from disk"""
        for profile_file in self.profiles_dir.glob("*.json"):
            try:
                with open(profile_file) as f:
                    data = json.load(f)
                    speaker = Speaker.from_dict(data)
                    self.speakers[speaker.id] = speaker
                    logger.info(f"Loaded speaker profile: {speaker.name} ({speaker.role.value})")
            except Exception as e:
                logger.error(f"Error loading profile {profile_file}: {e}")

    def _save_profile(self, speaker: Speaker):
        """Save speaker profile to disk"""
        profile_path = self.profiles_dir / f"{speaker.id}.json"
        speaker.last_seen = datetime.now().isoformat()
        with open(profile_path, "w") as f:
            json.dump(speaker.to_dict(), f, indent=2)

    def register_speaker(
        self,
        name: str,
        role: SpeakerRole,
        voice_sample: Optional[bytes] = None
    ) -> Speaker:
        """Register a new speaker"""
        speaker_id = hashlib.md5(name.lower().encode()).hexdigest()[:12]

        speaker = Speaker(
            id=speaker_id,
            name=name,
            role=role,
            created_at=datetime.now().isoformat()
        )

        if voice_sample:
            # TODO: Extract voice embedding from sample
            # For now, we'll use Kit API for this
            pass

        self.speakers[speaker_id] = speaker
        self._save_profile(speaker)
        logger.info(f"Registered speaker: {name} as {role.value}")
        return speaker

    async def identify_speaker(self, audio_data: bytes) -> Tuple[Optional[Speaker], float]:
        """
        Identify speaker from audio data.
        Returns (speaker, confidence) or (None, 0.0) if unknown.
        """
        # Use Kit API for speaker embedding extraction
        try:
            async with httpx.AsyncClient() as client:
                # First, transcribe to get the text
                files = {"file": ("audio.wav", audio_data, "audio/wav")}
                resp = await client.post(
                    f"{KIT_API_URL}/transcribe",
                    files=files,
                    timeout=30
                )

                if resp.status_code == 200:
                    result = resp.json()
                    # Kit API might return speaker info if diarization is enabled
                    speaker_id = result.get("speaker_id")

                    if speaker_id and speaker_id in self.speakers:
                        return self.speakers[speaker_id], 0.9

                    # Fallback: return unknown
                    return None, 0.0

        except Exception as e:
            logger.error(f"Speaker identification error: {e}")

        return None, 0.0

    def get_speaker_by_name(self, name: str) -> Optional[Speaker]:
        """Get speaker by name (case-insensitive)"""
        name_lower = name.lower()
        for speaker in self.speakers.values():
            if speaker.name.lower() == name_lower:
                return speaker
        return None


class TTSEngine:
    """
    Text-to-Speech engine with multiple backends.
    Primary: Google Cloud TTS (WaveNet voices)
    Fallback: Piper (local)
    """

    def __init__(self):
        self.google_enabled = GOOGLE_TTS_ENABLED
        self._google_client = None

    async def synthesize(
        self,
        text: str,
        voice: str = "nl-NL-Standard-A",
        use_google: bool = True
    ) -> Optional[bytes]:
        """
        Synthesize speech from text.

        Args:
            text: Text to speak
            voice: Voice ID (Google or Piper)
            use_google: Use Google TTS if available

        Returns:
            Audio data as bytes (WAV format)
        """
        if use_google and self.google_enabled:
            audio = await self._google_tts(text, voice)
            if audio:
                return audio

        # Fallback to Kit API / Piper
        return await self._piper_tts(text)

    async def _google_tts(self, text: str, voice: str) -> Optional[bytes]:
        """Google Cloud TTS via Vertex AI"""
        try:
            # Use Google Cloud TTS REST API
            from google.cloud import texttospeech

            client = texttospeech.TextToSpeechClient()

            synthesis_input = texttospeech.SynthesisInput(text=text)

            # Determine language from voice ID
            language_code = voice.split("-")[0] + "-" + voice.split("-")[1] if "-" in voice else "nl-NL"

            voice_params = texttospeech.VoiceSelectionParams(
                language_code=language_code,
                name=voice
            )

            audio_config = texttospeech.AudioConfig(
                audio_encoding=texttospeech.AudioEncoding.LINEAR16,
                sample_rate_hertz=SAMPLE_RATE
            )

            response = client.synthesize_speech(
                input=synthesis_input,
                voice=voice_params,
                audio_config=audio_config
            )

            logger.info(f"Google TTS: synthesized {len(text)} chars")
            return response.audio_content

        except ImportError:
            logger.warning("Google Cloud TTS not installed, falling back to Piper")
            return None
        except Exception as e:
            logger.error(f"Google TTS error: {e}")
            return None

    async def _piper_tts(self, text: str) -> Optional[bytes]:
        """Local TTS via Kit API / Piper"""
        try:
            async with httpx.AsyncClient() as client:
                resp = await client.post(
                    f"{KIT_API_URL}/tts",
                    json={"text": text, "voice": "default"},
                    timeout=30
                )

                if resp.status_code == 200:
                    logger.info(f"Piper TTS: synthesized {len(text)} chars")
                    return resp.content

        except Exception as e:
            logger.error(f"Piper TTS error: {e}")

        return None


class IntentParser:
    """
    Parse voice commands to extract intent and entities.
    Uses OomLlama for understanding.
    """

    INTENT_PROMPT = """Je bent een intent parser voor OomLlama, de huisassistent.

Analyseer deze zin en geef een JSON response:
"{text}"

Mogelijke intents:
- reminder: "herinner me", "over X minuten", "schoenen aan", etc.
- timer: "zet timer", "X minuten", etc.
- question: vraag stellen
- command: actie uitvoeren (licht aan, deur open, etc.)
- chat: gewoon praten

Geef ALLEEN geldige JSON:
{{"intent": "intent_type", "entities": {{"time_minutes": 0, "message": "", "target": ""}}, "confidence": 0.9}}"""

    async def parse(self, text: str) -> VoiceCommand:
        """Parse voice command to extract intent"""
        try:
            async with httpx.AsyncClient() as client:
                prompt = self.INTENT_PROMPT.format(text=text)

                resp = await client.post(
                    f"{OLLAMA_URL}/api/generate",
                    json={
                        "model": "qwen2.5:3b",  # Fast model for parsing
                        "prompt": prompt,
                        "stream": False,
                        "format": "json"
                    },
                    timeout=30
                )

                if resp.status_code == 200:
                    result = resp.json()
                    response_text = result.get("response", "{}")

                    # Parse JSON from response
                    try:
                        parsed = json.loads(response_text)
                        return VoiceCommand(
                            raw_text=text,
                            speaker=None,
                            intent=parsed.get("intent", "chat"),
                            entities=parsed.get("entities", {}),
                            confidence=parsed.get("confidence", 0.5),
                            timestamp=datetime.now().isoformat()
                        )
                    except json.JSONDecodeError:
                        logger.warning(f"Could not parse intent JSON: {response_text}")

        except Exception as e:
            logger.error(f"Intent parsing error: {e}")

        # Fallback: simple keyword matching
        return self._fallback_parse(text)

    def _fallback_parse(self, text: str) -> VoiceCommand:
        """Simple keyword-based intent detection"""
        text_lower = text.lower()

        intent = "chat"
        entities = {}

        # Check for reminder/timer keywords
        reminder_keywords = ["herinner", "remind", "over", "minuten", "seconden", "timer", "wekker"]
        if any(kw in text_lower for kw in reminder_keywords):
            intent = "reminder"

            # Extract time
            import re
            time_match = re.search(r'(\d+)\s*(minuten?|min|seconden?|sec)', text_lower)
            if time_match:
                value = int(time_match.group(1))
                unit = time_match.group(2)
                if "sec" in unit:
                    entities["time_seconds"] = value
                else:
                    entities["time_minutes"] = value

            entities["message"] = text

        # Check for question
        elif text_lower.endswith("?") or text_lower.startswith(("wat", "wie", "waar", "wanneer", "hoe", "waarom")):
            intent = "question"

        # Check for commands
        command_keywords = ["zet", "doe", "open", "sluit", "aan", "uit", "start", "stop"]
        if any(kw in text_lower for kw in command_keywords):
            intent = "command"

        return VoiceCommand(
            raw_text=text,
            speaker=None,
            intent=intent,
            entities=entities,
            confidence=0.6,
            timestamp=datetime.now().isoformat()
        )


class OomLlamaVoice:
    """
    Main OomLlama Voice Assistant.

    The friendly house assistant that:
    - Listens for "Yo Llama" wake word
    - Recognizes who is speaking
    - Responds appropriately based on role
    - Sets real timers via TIBET Scheduler
    """

    def __init__(self):
        self.speaker_recognition = SpeakerRecognition()
        self.tts = TTSEngine()
        self.intent_parser = IntentParser()
        self.running = False
        self.callbacks: Dict[str, Callable] = {}

        # Ensure data directories exist
        DATA_DIR.mkdir(parents=True, exist_ok=True)
        AUDIO_CACHE_DIR.mkdir(parents=True, exist_ok=True)

        # Register default speakers if none exist
        if not self.speaker_recognition.speakers:
            self._register_default_speakers()

    def _register_default_speakers(self):
        """Register default speaker profiles"""
        self.speaker_recognition.register_speaker("Jasper", SpeakerRole.ADMIN)
        self.speaker_recognition.register_speaker("Storm", SpeakerRole.FAMILY)
        self.speaker_recognition.register_speaker("Lieven", SpeakerRole.FAMILY)
        self.speaker_recognition.register_speaker("Pa", SpeakerRole.FAMILY)
        self.speaker_recognition.register_speaker("Gast", SpeakerRole.GUEST)
        logger.info("Registered default speaker profiles")

    def detect_wake_word(self, text: str) -> bool:
        """Check if text contains the wake word"""
        text_lower = text.lower().strip()
        return any(variant in text_lower for variant in WAKE_WORD_VARIANTS)

    def strip_wake_word(self, text: str) -> str:
        """Remove wake word from text"""
        text_lower = text.lower()
        for variant in WAKE_WORD_VARIANTS:
            if variant in text_lower:
                # Find position and remove
                pos = text_lower.find(variant)
                text = text[:pos] + text[pos + len(variant):]
                break
        return text.strip().strip(",").strip()

    async def process_audio(self, audio_data: bytes) -> Optional[str]:
        """
        Process audio input: transcribe and identify speaker.

        Returns transcribed text or None if failed.
        """
        try:
            async with httpx.AsyncClient() as client:
                files = {"file": ("audio.wav", audio_data, "audio/wav")}
                resp = await client.post(
                    f"{KIT_API_URL}/transcribe",
                    files=files,
                    timeout=30
                )

                if resp.status_code == 200:
                    result = resp.json()
                    return result.get("text", "")

        except Exception as e:
            logger.error(f"Audio processing error: {e}")

        return None

    async def handle_command(
        self,
        text: str,
        speaker: Optional[Speaker] = None
    ) -> Dict:
        """
        Handle a voice command.

        This is the main entry point for processing commands.
        """
        # Parse intent
        command = await self.intent_parser.parse(text)
        command.speaker = speaker

        logger.info(f"Command: '{text}' -> intent={command.intent}, speaker={speaker.name if speaker else 'unknown'}")

        # Check permissions
        if not self._check_permission(command):
            return await self._respond_no_permission(command)

        # Route by intent
        if command.intent == "reminder":
            return await self._handle_reminder(command)
        elif command.intent == "timer":
            return await self._handle_timer(command)
        elif command.intent == "question":
            return await self._handle_question(command)
        elif command.intent == "command":
            return await self._handle_device_command(command)
        else:
            return await self._handle_chat(command)

    def _check_permission(self, command: VoiceCommand) -> bool:
        """Check if speaker has permission for this command"""
        if not command.speaker:
            # Unknown speaker - allow basic commands only
            return command.intent in ["chat", "question"]

        role = command.speaker.role

        # Admin can do everything
        if role == SpeakerRole.ADMIN:
            return True

        # Family can do most things
        if role == SpeakerRole.FAMILY:
            return command.intent in ["reminder", "timer", "question", "chat"]

        # Guest has limited access
        if role == SpeakerRole.GUEST:
            return command.intent in ["question", "chat"]

        # Medical has special access
        if role == SpeakerRole.MEDICAL:
            return True

        return False

    async def _respond_no_permission(self, command: VoiceCommand) -> Dict:
        """Generate response for permission denied"""
        message = "Sorry, daar heb je geen toestemming voor."

        if command.speaker and command.speaker.role == SpeakerRole.FAMILY:
            message = f"Sorry {command.speaker.name}, vraag dat even aan papa."

        return {
            "success": False,
            "intent": command.intent,
            "response": message,
            "speak": True
        }

    async def _handle_reminder(self, command: VoiceCommand) -> Dict:
        """Handle reminder intent - schedule via TIBET"""
        entities = command.entities
        minutes = entities.get("time_minutes", 0)
        seconds = entities.get("time_seconds", 0)
        message = entities.get("message", command.raw_text)

        # Clean up message
        message = self.strip_wake_word(message)

        # Default to 5 minutes if no time specified
        if minutes == 0 and seconds == 0:
            minutes = 5

        try:
            async with httpx.AsyncClient() as client:
                resp = await client.post(
                    f"{TIBET_SCHEDULER_URL}/remind",
                    json={
                        "message": message,
                        "minutes": minutes,
                        "seconds": seconds,
                        "reason": f"Voice reminder from {command.speaker.name if command.speaker else 'unknown'}"
                    },
                    timeout=10
                )

                if resp.status_code == 200:
                    result = resp.json()

                    # Generate friendly response
                    time_str = ""
                    if minutes > 0:
                        time_str += f"{minutes} minuten"
                    if seconds > 0:
                        if time_str:
                            time_str += " en "
                        time_str += f"{seconds} seconden"

                    response = f"Oké, ik herinner je over {time_str}!"

                    return {
                        "success": True,
                        "intent": "reminder",
                        "scheduled_id": result.get("scheduled_id"),
                        "fires_in": result.get("fires_in"),
                        "response": response,
                        "speak": True
                    }

        except Exception as e:
            logger.error(f"Reminder scheduling error: {e}")

        return {
            "success": False,
            "intent": "reminder",
            "response": "Sorry, ik kon de herinnering niet instellen.",
            "speak": True
        }

    async def _handle_timer(self, command: VoiceCommand) -> Dict:
        """Handle timer intent - same as reminder but different wording"""
        return await self._handle_reminder(command)

    async def _handle_question(self, command: VoiceCommand) -> Dict:
        """Handle question - ask OomLlama"""
        return await self._ask_oomllama(command.raw_text, command.speaker)

    async def _handle_chat(self, command: VoiceCommand) -> Dict:
        """Handle general chat - talk to OomLlama"""
        return await self._ask_oomllama(command.raw_text, command.speaker)

    async def _handle_device_command(self, command: VoiceCommand) -> Dict:
        """Handle device commands (lights, doors, etc.)"""
        # TODO: Integrate with home automation
        return {
            "success": False,
            "intent": "command",
            "response": "Apparaatbesturing is nog niet geïmplementeerd.",
            "speak": True
        }

    async def _ask_oomllama(
        self,
        question: str,
        speaker: Optional[Speaker] = None
    ) -> Dict:
        """Ask OomLlama a question"""

        # Build context-aware prompt
        speaker_context = ""
        if speaker:
            speaker_context = f"\nDe gebruiker is {speaker.name} ({speaker.role.value})."
            if speaker.role == SpeakerRole.FAMILY and speaker.name in ["Storm", "Lieven"]:
                speaker_context += " Dit is een kind, geef kindvriendelijke antwoorden."

        prompt = f"""Je bent OomLlama, de vriendelijke huisassistent van de familie Van de Meent.
Je bent warm, behulpzaam en hebt een vleugje humor.
{speaker_context}

Vraag: {question}

Geef een kort, vriendelijk antwoord in het Nederlands (max 3 zinnen)."""

        try:
            async with httpx.AsyncClient() as client:
                resp = await client.post(
                    f"{OLLAMA_URL}/api/generate",
                    json={
                        "model": "qwen2.5:7b",
                        "prompt": prompt,
                        "stream": False
                    },
                    timeout=60
                )

                if resp.status_code == 200:
                    result = resp.json()
                    response = result.get("response", "").strip()

                    return {
                        "success": True,
                        "intent": "chat",
                        "response": response,
                        "speak": True
                    }

        except Exception as e:
            logger.error(f"OomLlama error: {e}")

        return {
            "success": False,
            "intent": "chat",
            "response": "Sorry, ik kon even niet denken. Probeer het nog eens.",
            "speak": True
        }

    async def speak(self, text: str, use_google: bool = True) -> bool:
        """Speak text using TTS"""
        audio = await self.tts.synthesize(text, use_google=use_google)

        if audio:
            # TODO: Play audio via Kit API or local audio
            # For now, log it
            logger.info(f"Speaking: {text}")

            # Send to Kit API for playback
            try:
                async with httpx.AsyncClient() as client:
                    await client.post(
                        f"{KIT_API_URL}/play",
                        files={"file": ("speech.wav", audio, "audio/wav")},
                        timeout=30
                    )
                return True
            except Exception as e:
                logger.error(f"Audio playback error: {e}")

        return False

    async def process_and_respond(
        self,
        text: str,
        speaker_name: Optional[str] = None,
        speak_response: bool = True
    ) -> Dict:
        """
        Process text command and optionally speak the response.

        This is the main API for external callers.
        """
        # Get speaker if name provided
        speaker = None
        if speaker_name:
            speaker = self.speaker_recognition.get_speaker_by_name(speaker_name)

        # Check for wake word
        has_wake_word = self.detect_wake_word(text)

        if has_wake_word:
            text = self.strip_wake_word(text)

        # Handle the command
        result = await self.handle_command(text, speaker)

        # Speak if requested
        if speak_response and result.get("speak"):
            await self.speak(result.get("response", ""))

        return result


# ============================================================================
# REST API Integration
# ============================================================================

async def handle_voice_request(
    text: str,
    speaker: str = None,
    audio_base64: str = None
) -> Dict:
    """
    Handle a voice request (called from API).

    Can receive either text or base64 audio.
    """
    voice = OomLlamaVoice()

    if audio_base64:
        # Decode and transcribe audio
        import base64
        audio_data = base64.b64decode(audio_base64)
        text = await voice.process_audio(audio_data)

        if not text:
            return {"success": False, "error": "Could not transcribe audio"}

    return await voice.process_and_respond(text, speaker_name=speaker)


# ============================================================================
# Main Entry Point
# ============================================================================

async def main():
    """Interactive test mode"""
    print("""
    ╔════════════════════════════════════════════════════════════╗
    ║           OOMLLAMA VOICE MODULE                            ║
    ║           The friendly house assistant                     ║
    ║                                                            ║
    ║   Wake word: "Yo Llama"                                    ║
    ║   Type commands to test, or 'quit' to exit                 ║
    ║                                                            ║
    ║   Voor pa. Voor iedereen die begrepen wil worden.          ║
    ║                                                            ║
    ║   One love, one fAmIly                                     ║
    ╚════════════════════════════════════════════════════════════╝
    """)

    voice = OomLlamaVoice()

    while True:
        try:
            text = input("\n🎤 Jij: ").strip()

            if text.lower() in ["quit", "exit", "stop"]:
                print("Tot ziens!")
                break

            if not text:
                continue

            # Test as Jasper (admin)
            result = await voice.process_and_respond(
                text,
                speaker_name="Jasper",
                speak_response=False  # Don't try to play audio in CLI
            )

            print(f"🦙 OomLlama: {result.get('response', 'No response')}")
            print(f"   [intent: {result.get('intent')}, success: {result.get('success')}]")

        except KeyboardInterrupt:
            print("\n\nTot ziens!")
            break
        except Exception as e:
            print(f"Error: {e}")


if __name__ == "__main__":
    asyncio.run(main())
