import asyncio
import os
import time
import traceback
from pathlib import Path

import numpy as np
import pyaudio
import webrtcvad
from openwakeword.model import Model
from faster_whisper import WhisperModel
from edge_tts import Communicate

from agent import Agent
from util.logger import get_logger

logger = get_logger("NORA")


# =========================
# CONFIG
# =========================
WAKEWORD = "nora"
WAKE_THRESHOLD = 0.60

SAMPLE_RATE = 16000
FRAME_SIZE = 1024

VAD_MODE = 3  # 0-3 (3 = aggressiv)
MAX_RECORD_SEC = 12
SILENCE_LIMIT = 25  # frames

TTS_VOICE = "de-DE-AmalaNeural"

MODEL_PATH = Path.home() / ".cache/openwakeword/models/hey_jarvis_v0.1.tflite"


# =========================
# VOICE ASSISTANT CORE
# =========================
class VoiceNora:
    def __init__(self):
        logger.info("Initialisiere N.O.R.A...")

        self.agent = Agent()
        self.running = True

        # ---- VAD ----
        self.vad = webrtcvad.Vad(VAD_MODE)

        # ---- Wakeword Model ----
        self.ww_model = self._load_wakeword_model()

        # ---- Whisper STT ----
        self.stt = WhisperModel(
            "tiny",
            device="cpu",
            compute_type="int8"
        )

        logger.info("N.O.R.A bereit ✔")
        logger.info(f"Wakeword: {WAKEWORD}")
        logger.info(f"Model: {MODEL_PATH}")

    # =========================
    # INIT HELPERS
    # =========================
    def _load_wakeword_model(self):
        logger.info("Lade Wakeword-Modell...")

        if not MODEL_PATH.exists():
            logger.error(f"Model nicht gefunden: {MODEL_PATH}")
            raise FileNotFoundError(MODEL_PATH)

        try:
            model = Model(
                wakeword_models=[str(MODEL_PATH)]
            )
            logger.info("Wakeword-Modell geladen ✔")
            return model

        except Exception as e:
            logger.error("Fehler beim Laden des Wakeword-Modells")
            logger.error(str(e))
            raise

    # =========================
    # AUDIO STREAM
    # =========================
    def _open_stream(self):
        pa = pyaudio.PyAudio()
        stream = pa.open(
            format=pyaudio.paInt16,
            channels=1,
            rate=SAMPLE_RATE,
            input=True,
            frames_per_buffer=FRAME_SIZE
        )
        return pa, stream

    # =========================
    # MAIN LOOP
    # =========================
    async def run(self):
        logger.info("Starte Audio-Loop...")

        pa, stream = self._open_stream()

        try:
            while self.running:
                audio = stream.read(FRAME_SIZE, exception_on_overflow=False)
                audio_np = np.frombuffer(audio, dtype=np.int16)

                try:
                    result = self.ww_model.predict(audio_np)

                    score = max(result.values()) if result else 0.0

                    if score > WAKE_THRESHOLD:
                        logger.info(f"Wakeword erkannt (score={score:.2f})")
                        await self.handle_command()

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

        except KeyboardInterrupt:
            logger.info("Beende N.O.R.A...")

        finally:
            stream.stop_stream()
            stream.close()
            pa.terminate()

    # =========================
    # COMMAND FLOW
    # =========================
    async def handle_command(self):
        logger.info("Aufnahme startet...")

        audio = await self.record_until_silence()

        if len(audio) < 8000:
            logger.warning("Audio zu kurz")
            return

        text = await self.transcribe(audio)

        if not text:
            logger.warning("Keine Sprache erkannt")
            return

        logger.info(f"USER: {text}")

        response = await self.agent.run(text)

        logger.info(f"NORA: {response}")

        await self.speak(response)

    # =========================
    # RECORDING (VAD)
    # =========================
    async def record_until_silence(self):
        pa, stream = self._open_stream()

        frames = []
        silence = 0

        max_frames = int(SAMPLE_RATE / FRAME_SIZE * MAX_RECORD_SEC)

        for _ in range(max_frames):
            chunk = stream.read(FRAME_SIZE, exception_on_overflow=False)
            frames.append(chunk)

            is_speech = self.vad.is_speech(chunk, SAMPLE_RATE)

            if is_speech:
                silence = 0
            else:
                silence += 1

            if silence > SILENCE_LIMIT and len(frames) > 20:
                break

        stream.stop_stream()
        stream.close()
        pa.terminate()

        return b"".join(frames)

    # =========================
    # STT
    # =========================
    async def transcribe(self, audio_bytes):
        try:
            segments, _ = self.stt.transcribe(
                audio_bytes,
                language="de",
                beam_size=5
            )

            text = " ".join([s.text for s in segments]).strip()
            return text

        except Exception as e:
            logger.error("STT Fehler")
            logger.error(traceback.format_exc())
            return ""

    # =========================
    # TTS
    # =========================
    async def speak(self, text: str):
        try:
            filename = f"response_{int(time.time())}.mp3"

            tts = Communicate(text, voice=TTS_VOICE)
            await tts.save(filename)

            os.system(f"mpg123 -q {filename}")

            os.remove(filename)

        except Exception as e:
            logger.error("TTS Fehler")
            logger.error(traceback.format_exc())


# =========================
# MAIN
# =========================
async def main():
    nora = VoiceNora()
    await nora.run()


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