# app/main.py

from fastapi import FastAPI
from fastapi.responses import JSONResponse
from aiortc import RTCPeerConnection, RTCSessionDescription
from aiortc.contrib.media import MediaStreamTrack
from app.ai_brain import PersonaPlexBrain
import numpy as np

app = FastAPI()
brain = PersonaPlexBrain()


class AIStreamTrack(MediaStreamTrack):
    kind = "audio"

    def __init__(self, source_track):
        super().__init__()
        self.source = source_track

    async def recv(self):
        frame = await self.source.recv()

        audio = frame.to_ndarray()[0]
        out = brain.process_audio_frame(audio)

        if out is None:
            return frame

        frame.planes[0].update(out.tobytes())
        return frame


@app.post("/offer")
async def offer(offer: dict):
    pc = RTCPeerConnection()

    @pc.on("track")
    def on_track(track):
        if track.kind == "audio":
            pc.addTrack(AIStreamTrack(track))

    await pc.setRemoteDescription(
        RTCSessionDescription(sdp=offer["sdp"], type=offer["type"])
    )

    answer = await pc.createAnswer()
    await pc.setLocalDescription(answer)

    return JSONResponse(
        {"sdp": pc.localDescription.sdp, "type": pc.localDescription.type}
    )