server.py

from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Request
from fastapi.responses import JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from aiortc import RTCPeerConnection, RTCSessionDescription
import cv2
import wave
import io
import numpy as np
from collections import deque
from aiortc.contrib.media import MediaStreamTrack
import asyncio
import json
import logging
import fractions

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI()

# Enable CORS for cross-origin requests
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Pydantic model for WebRTC offer
class RTCOffer(BaseModel):
    sdp: str
    type: str

# Store received frames and audio
received_frames = deque(maxlen=20)
received_audio_chunks = deque(maxlen=1000)
stream_stats = {"video_frames": 0, "audio_chunks": 0, "connected": False}
peer_connections = set()  # Keep track of active peer connections

@app.get("/")
def home():
    return {
        "status": "Server running",
        "stats": stream_stats,
        "frames_received": len(received_frames)
    }

@app.post("/offer")
async def handle_offer(offer_data: RTCOffer):
    """Handle WebRTC offer from Raspberry Pi"""
    try:
        logger.info("📨 Received WebRTC offer from client")
        offer = RTCSessionDescription(sdp=offer_data.sdp, type=offer_data.type)
        
        # Create peer connection
        pc = RTCPeerConnection()
        peer_connections.add(pc)
        stream_stats["connected"] = True
        
        @pc.on("track")
        async def on_track(track):
            logger.info(f"✓ Track received: {track.kind}")
            stream_stats[f"{track.kind}_frames"] = 0

            counter = 0
            while True:
                try:
                    frame = await track.recv()
                    if track.kind == "video":
                        counter+=1
                        frame.to_image().save(f"images/my_img{counter}.jpg")
                        
                        stream_stats["video_frames"] += 1
                        received_frames.append({
                            "timestamp": asyncio.get_event_loop().time(),
                            "data": frame.to_ndarray(format="bgr24")  # Convert to numpy
                        })
                        logger.info(f"📹 Frame #{stream_stats['video_frames']}")
                        
                    elif track.kind == "audio":
                        stream_stats["audio_chunks"] += 1
                        received_audio_chunks.append({
                            "pts": frame.pts,
                            "timestamp": asyncio.get_event_loop().time(),
                            "data": frame.to_ndarray().tobytes()

                        })
                        logger.info(f"🎤 Audio chunk #{stream_stats['audio_chunks']} received")

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

        
        # IMPORTANT: Set remote description BEFORE creating answer
        await pc.setRemoteDescription(offer)
        
        # Create answer
        answer = await pc.createAnswer()
        await pc.setLocalDescription(answer)
        
        logger.info("✓ WebRTC connection established")
        return {
            "sdp": pc.localDescription.sdp,
            "type": pc.localDescription.type
        }
    
    except Exception as e:
        import traceback
        logger.error(f"❌ Error handling offer: {type(e).__name__}: {e}")
        traceback.print_exc()
        return {"error": str(e), "type": type(e).__name__}, 500

@app.get("/stream/status")
async def stream_status():
    """Get current streaming status"""
    return {
        "stats": stream_stats,
        "video_frames_received": len(received_frames),
        "audio_chunks_received": len(received_audio_chunks),
        "latest_frames": len(received_frames[-5:]) if received_frames else 0
    }

@app.get("/stream/get-video-frame/{index}")
async def get_video_frame(index: int):
    """Retrieve a specific video frame (returns bytes)"""
    if 0 <= index < len(received_frames):
        frame_data = received_frames[index]["data"]
        return JSONResponse({
            "index": index,
            "timestamp": received_frames[index]["timestamp"],
            "data_type": str(type(frame_data))
        })
    return {"error": "Frame not found"}, 404

@app.post("/audio/save-wav")
async def save_audio_wav():
    """Reconstruct and save all buffered audio to WAV file"""
    if not received_audio_chunks:
        return {"error": "No audio data"}
    
    sorted_frames = sorted(received_audio_chunks, key=lambda x: x["pts"]) 
    all_data = b''.join(frame["data"] for frame in sorted_frames)
    
    # Create WAV in memory
    with wave.open("mywav.wav", 'wb') as wf:
        wf.setnchannels(1)  # mono
        wf.setsampwidth(2)  # 16-bit
        wf.setframerate(48000)
        wf.writeframes(all_data)
    
    return {"succcess": "file saved"}

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=5000, log_level="info")