from fastapi import FastAPI, WebSocket, WebSocketDisconnect, UploadFile, File, Request
from fastapi.responses import HTMLResponse
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
import os
import shutil
import uuid
import uvicorn

app = FastAPI()

# =========================
# 基础配置
# =========================

BASE_DIR = os.path.dirname(os.path.abspath(__file__))

TEMPLATE_DIR = os.path.join(BASE_DIR, "templates")
UPLOAD_DIR = os.path.join(BASE_DIR, "uploads")

MAX_FILE_SIZE = 10 * 1024 * 1024  # 10MB

ALLOWED_IMAGE_EXTENSIONS = {"png", "jpg", "jpeg", "gif"}
ALLOWED_VIDEO_EXTENSIONS = {"mp4", "webm"}

# =========================
# 目录初始化
# =========================

os.makedirs(UPLOAD_DIR, exist_ok=True)

templates = Jinja2Templates(directory=TEMPLATE_DIR)
app.mount("/uploads", StaticFiles(directory=UPLOAD_DIR), name="uploads")

# =========================
# WebSocket 状态
# =========================

active_connections: list[WebSocket] = []
usernames: dict[WebSocket, str] = {}

# =========================
# 工具函数
# =========================

def get_ext(filename: str) -> str:
    return filename.rsplit(".", 1)[-1].lower()

def is_image(ext: str) -> bool:
    return ext in ALLOWED_IMAGE_EXTENSIONS

def is_video(ext: str) -> bool:
    return ext in ALLOWED_VIDEO_EXTENSIONS

# =========================
# 页面路由
# =========================

@app.get("/", response_class=HTMLResponse)
async def index(request: Request):
    return templates.TemplateResponse("chat.html", {"request": request})

# =========================
# 文件上传接口
# =========================

@app.post("/upload")
async def upload_file(file: UploadFile = File(...)):
    ext = get_ext(file.filename)

    if not (is_image(ext) or is_video(ext)):
        return {"error": "不支持的文件类型"}

    contents = await file.read()

    if len(contents) > MAX_FILE_SIZE:
        return {"error": "文件过大（最大 10MB）"}

    filename = f"{uuid.uuid4().hex}.{ext}"
    save_path = os.path.join(UPLOAD_DIR, filename)

    with open(save_path, "wb") as f:
        f.write(contents)

    file_url = f"/uploads/{filename}"

    if is_image(ext):
        msg = f"IMG::{file_url}"
    elif is_video(ext):
        msg = f"VIDEO::{file_url}"
    else:
        msg = f"FILE::{file_url}"

    await broadcast(msg)

    return {"url": file_url}

# =========================
# WebSocket 聊天
# =========================

@app.websocket("/ws")
async def websocket_endpoint(ws: WebSocket):
    await ws.accept()
    active_connections.append(ws)
    await broadcast_online_count()

    try:
        # 首条消息作为用户名
        username = await ws.receive_text()
        usernames[ws] = username

        await broadcast(f"SYSTEM::{username} 加入聊天室")
        await broadcast_online_count()

        while True:
            text = await ws.receive_text()
            await broadcast(f"TEXT::{username}::{text}")

    except WebSocketDisconnect:
        pass

    finally:
        if ws in active_connections:
            active_connections.remove(ws)
        name = usernames.pop(ws, "匿名用户")
        await broadcast(f"SYSTEM::{name} 离开聊天室")
        await broadcast_online_count()

# =========================
# 广播工具
# =========================

async def broadcast(message: str):
    dead = []
    for ws in active_connections:
        try:
            await ws.send_text(message)
        except Exception:
            dead.append(ws)

    for ws in dead:
        active_connections.remove(ws)
        usernames.pop(ws, None)

async def broadcast_online_count():
    msg = f"COUNT::{len(active_connections)}"
    await broadcast(msg)

# =========================
# 启动
# =========================

if __name__ == "__main__":
    uvicorn.run(
        "chat:app",
        host="0.0.0.0",
        port=8202,
        reload=True
    )