import io
import numpy as np
import soundfile as sf
import edge_tts
import miniaudio

from fastapi import FastAPI
from fastapi.responses import StreamingResponse
from scipy.signal import resample
from pydantic import BaseModel

from ltm import EmbeddingMemory

app = FastAPI()
memory = EmbeddingMemory()

VOICE = "pl-PL-ZofiaNeural"

class MemoryMessage(BaseModel):
    session_id: str
    role: str
    content: str

class MemoryQuery(BaseModel):
    session_id: str | None = None
    query: str
    top_k: int = 5

@app.post("/tts")
async def tts(data: dict):
    text = data.get("text", "").strip()
    if not text:
        return {"error": "empty text"}

    mp3_buffer = io.BytesIO()
    communicate = edge_tts.Communicate(text, VOICE)

    async for chunk in communicate.stream():
        if chunk["type"] == "audio":
            mp3_buffer.write(chunk["data"])

    mp3_buffer.seek(0)

    decoded = miniaudio.decode(mp3_buffer.read())

    raw = decoded.samples

    if decoded.sample_format == miniaudio.SampleFormat.SIGNED16:
        audio = np.frombuffer(raw, dtype=np.int16).astype(np.float32) / 32768.0

    elif decoded.sample_format == miniaudio.SampleFormat.FLOAT32:
        audio = np.frombuffer(raw, dtype=np.float32)

    else:
        return {"error": "Unsupported audio format from decoder"}

    if decoded.nchannels > 1:
        audio = audio.reshape(-1, decoded.nchannels)
        audio = np.mean(audio, axis=1)

    peak = np.max(np.abs(audio)) if len(audio) else 0
    if peak > 0:
        audio = audio / max(peak, 1.0)

    audio *= 0.3

    pitch_factor = 1.1
    audio = resample(audio, int(len(audio) / pitch_factor))

    audio = np.clip(audio, -1.0, 1.0)

    wav_buffer = io.BytesIO()
    sf.write(wav_buffer, audio, decoded.sample_rate, format="WAV")
    wav_buffer.seek(0)

    return StreamingResponse(wav_buffer, media_type="audio/wav")

@app.post("/memory/store")
def store(msg: MemoryMessage):
    memory.add_message(
        session_id=msg.session_id,
        role=msg.role,
        content=msg.content
    )

    return {"status": "saved"}

@app.post("/memory/query")
def search_memory(data: MemoryQuery):
    results = memory.search(
        query=data.query,
        session_id=data.session_id,
        top_k=data.top_k
    )

    return results