import asyncio
import subprocess
import websockets
from fastapi import FastAPI, WebSocket, WebSocketDisconnect

app = FastAPI()

# 管理推拉流的 ffmpeg 子进程
processes = {}

def start_stream(input_url: str, output_url: str):
    command = [
        "ffmpeg",
        "-i", input_url,
        "-c:v", "libx264",
        "-c:a", "aac",
        "-f", "flv",
        output_url
    ]
    process = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    processes[output_url] = process
    return process

def stop_stream(output_url: str):
    process = processes.pop(output_url, None)
    if process:
        process.terminate()

# WebSocket 心跳检测
@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    try:
        while True:
            message = await websocket.receive_text()
            if message == "ping":
                await websocket.send_text("pong")
    except WebSocketDisconnect:
        print("WebSocket disconnected")

# 启动推流示例
@app.post("/start-stream")
def start(input_url: str, output_url: str):
    process = start_stream(input_url, output_url)
    return {"message": "Stream started", "pid": process.pid}

# 停止推流示例
@app.post("/stop-stream")
def stop(output_url: str):
    stop_stream(output_url)
    return {"message": "Stream stopped"}