# voice_server.py
import os
import io
import asyncio
import json
import tempfile
from pathlib import Path
from dotenv import load_dotenv
import numpy as np
import soundfile as sf
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.staticfiles import StaticFiles
from openai import OpenAI
from anthropic import Anthropic
from elevenlabs import generate, set_api_key, Voice, VoiceSettings

load_dotenv()
set_api_key(os.environ["ELEVENLABS_API_KEY"])

app = FastAPI()
oai = OpenAI(api_key=os.environ["OPENAI_API_KEY"])
claude = Anthropic(api_key=os.environ["ANTHROPIC_API_KEY"])

VOICE_ID = os.environ["ELEVENLABS_VOICE_ID"]

SYSTEM_PROMPT = """You are a helpful voice assistant. Your responses will be
converted to speech, so follow these rules:

- Keep responses concise: 1-3 sentences for simple questions, 4-6 for complex
- Never use markdown, bullet points, or formatting — just natural speech
- Spell out abbreviations (e.g., say "kilobytes" not "KB")
- Avoid starting responses with filler phrases like "Certainly!" or "Of course!"
- If you need to list items, use natural connectors: "First..., then..., and finally..."
- Use a conversational tone"""

class VoiceSession:
    def __init__(self, session_id: str):
        self.session_id = session_id
        self.history: list[dict] = []

    def add_turn(self, user: str, assistant: str):
        self.history.append({"role": "user", "content": user})
        self.history.append({"role": "assistant", "content": assistant})
        # Keep last 10 turns (20 messages)
        if len(self.history) > 20:
            self.history = self.history[-20:]

sessions: dict[str, VoiceSession] = {}

def transcribe_audio(audio_bytes: bytes) -> str:
    """Transcribe audio using Whisper."""
    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
        f.write(audio_bytes)
        f.flush()
        with open(f.name, "rb") as audio_file:
            transcript = oai.audio.transcriptions.create(
                model="whisper-1",
                file=audio_file,
                language="en"
            )
    return transcript.text

def generate_response(text: str, session: VoiceSession) -> str:
    """Generate a response using Claude."""
    messages = session.history + [{"role": "user", "content": text}]

    response = claude.messages.create(
        model="claude-opus-4-6",
        max_tokens=512,
        system=SYSTEM_PROMPT,
        messages=messages
    )
    return response.content[0].text

def text_to_speech(text: str) -> bytes:
    """Convert text to speech using ElevenLabs."""
    audio = generate(
        text=text,
        voice=Voice(
            voice_id=VOICE_ID,
            settings=VoiceSettings(
                stability=0.5,
                similarity_boost=0.75,
                style=0.0,
                use_speaker_boost=True
            )
        ),
        model="eleven_multilingual_v2"
    )
    return audio

@app.websocket("/ws/{session_id}")
async def voice_websocket(websocket: WebSocket, session_id: str):
    await websocket.accept()

    if session_id not in sessions:
        sessions[session_id] = VoiceSession(session_id)
    session = sessions[session_id]

    try:
        while True:
            # Receive audio chunk from client
            data = await websocket.receive_bytes()

            # Signal: processing started
            await websocket.send_json({"type": "processing", "stage": "transcribing"})

            # Step 1: Transcribe
            try:
                transcript = await asyncio.get_event_loop().run_in_executor(
                    None, transcribe_audio, data
                )
            except Exception as e:
                await websocket.send_json({"type": "error", "message": str(e)})
                continue

            if not transcript.strip():
                await websocket.send_json({"type": "silence"})
                continue

            await websocket.send_json({
                "type": "transcript",
                "text": transcript
            })

            # Step 2: Generate response
            await websocket.send_json({"type": "processing", "stage": "thinking"})
            response_text = await asyncio.get_event_loop().run_in_executor(
                None, generate_response, transcript, session
            )

            await websocket.send_json({
                "type": "response_text",
                "text": response_text
            })

            # Step 3: Convert to speech
            await websocket.send_json({"type": "processing", "stage": "speaking"})
            audio_bytes = await asyncio.get_event_loop().run_in_executor(
                None, text_to_speech, response_text
            )

            # Update session history
            session.add_turn(transcript, response_text)

            # Send audio back
            await websocket.send_bytes(audio_bytes)
            await websocket.send_json({"type": "done"})

    except WebSocketDisconnect:
        pass
    except Exception as e:
        await websocket.send_json({"type": "error", "message": str(e)})

import re

async def stream_response_as_audio(
    websocket: WebSocket,
    response_text: str
):
    """Split response into sentences, TTS each sentence separately."""
    sentences = re.split(r'(?<=[.!?])\s+', response_text)

    for sentence in sentences:
        if not sentence.strip():
            continue
        audio = await asyncio.get_event_loop().run_in_executor(
            None, text_to_speech, sentence
        )
        await websocket.send_bytes(audio)
        await websocket.send_json({"type": "sentence_done"})

    await websocket.send_json({"type": "done"})