"""
ChatHub - 实时聊天室
FastAPI + WebSocket
"""
from fastapi import FastAPI, WebSocket, WebSocketDisconnect, UploadFile, File, Form, Request
from fastapi.responses import HTMLResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
from typing import List, Dict, Tuple
import os
import uuid
import json
import asyncio
import shutil
import time
import collections
from datetime import datetime

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")
BACKUP_DIR = os.path.join(BASE_DIR, "backups")
os.makedirs(UPLOAD_DIR, exist_ok=True)
os.makedirs(BACKUP_DIR, exist_ok=True)

CHAT_LOG_FILE = os.path.join(BASE_DIR, "chat_history.txt")
MAX_FILE_SIZE = 10 * 1024 * 1024  # 10MB per file

ALLOWED_IMAGE_EXTENSIONS = {"png", "jpg", "jpeg", "gif", "bmp", "tiff", "webp", "heic"}
ALLOWED_VIDEO_EXTENSIONS = {"mp4", "mov", "wmv", "avi", "m4v", "mpg", "mpeg", "flv", "mkv", "3gp", "webm"}

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

# ── Global State ────────────────────────────────────────────
connections: Dict[WebSocket, str] = {}  # ws -> username
glock = asyncio.Lock()

# ── Rate Limiting ─────────────────────────────────────────
# Upload rate limit: {ip: [timestamp, ...]}
upload_history: Dict[str, List[float]] = {}
UPLOAD_COOLDOWN = 3.0  # seconds between uploads

# Message rate limit: per-connection message timestamps
msg_history: Dict[WebSocket, List[float]] = {}
MSG_BURST = 8       # max messages
MSG_WINDOW = 3.0    # per this many seconds


def get_ext(fn: str) -> str:
    return fn.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


def log_to_file(msg: str):
    ts = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    try:
        with open(CHAT_LOG_FILE, "a", encoding="utf-8") as f:
            f.write(f"[{ts}] {msg}\n")
    except Exception:
        pass


def load_recent_history(n: int = 200) -> List[str]:
    if not os.path.exists(CHAT_LOG_FILE):
        return []
    try:
        with open(CHAT_LOG_FILE, "r", encoding="utf-8") as f:
            lines = f.readlines()
        return [l.strip() for l in lines[-n:]]
    except Exception:
        return []


# ── HTTP Endpoints ──────────────────────────────────────────

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


@app.get("/members")
async def get_members():
    """Return current online member list."""
    async with glock:
        names = list(connections.values())
    return JSONResponse({"members": names})


# ── Rate Limiter Helpers ────────────────────────────────

async def _check_upload_rate(request: Request) -> Tuple[bool, str]:
    """Check upload rate limit per client IP.
    Returns (allowed, reason).
    """
    ip = request.client.host if request.client else "unknown"
    now = time.time()
    
    async with glock:
        times = upload_history.get(ip, [])
        # Clean old entries
        times = [t for t in times if now - t < UPLOAD_COOLDOWN]
        if times:
            remaining = int(UPLOAD_COOLDOWN - (now - times[-1]))
            return False, f"上传太频繁，请 {remaining} 秒后再试"
        times.append(now)
        upload_history[ip] = times
    return True, ""


async def _check_msg_rate(ws: WebSocket) -> bool:
    """Check message send rate per WebSocket connection."""
    now = time.time()
    async with glock:
        times = msg_history.get(ws, [])
        times = [t for t in times if now - t < MSG_WINDOW]
        if len(times) >= MSG_BURST:
            return False
        times.append(now)
        msg_history[ws] = times
    return True


# ── HTTP Endpoints ──────────────────────────────────────────

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


@app.get("/members")
async def get_members():
    """Return current online member list."""
    async with glock:
        names = list(connections.values())
    return JSONResponse({"members": names})


@app.post("/upload")
async def upload_file(file: UploadFile = File(...), request: Request = None, sender: str = Form("匿名")):
    # Rate limit check
    allowed, reason = await _check_upload_rate(request)
    if not allowed:
        return JSONResponse({"error": reason}, status_code=429)
    
    # File type check
    ext = get_ext(file.filename)
    if not (is_image(ext) or is_video(ext)):
        return JSONResponse({"error": "不支持的文件类型，仅接受图片和视频"}, status_code=400)
    
    # Read and check size
    data = await file.read()
    if len(data) > MAX_FILE_SIZE:
        size_mb = len(data) / (1024 * 1024)
        return JSONResponse({
            "error": f"文件过大（{size_mb:.1f}MB），最大允许 10MB"
        }, status_code=413)
    
    # Save and broadcast
    fn = f"{uuid.uuid4().hex}.{ext}"
    path = os.path.join(UPLOAD_DIR, fn)
    with open(path, "wb") as f:
        f.write(data)
    url = f"/uploads/{fn}"
    # Format: TYPE::sender::url — frontend parses two colons
    msg = f"IMG::{sender}::{url}" if is_image(ext) else f"VIDEO::{sender}::{url}"
    asyncio.create_task(_broadcast(msg))
    return {"url": url}


# ── Broadcast ───────────────────────────────────────────────

async def _broadcast(message: str):
    """Send message to all connected clients.
    Uses snapshot-then-send pattern: grab targets under lock, then send concurrently without lock.
    """
    async with glock:
        targets = list(connections.items())  # [(ws, name), ...]

    tasks = []
    for ws, _ in targets:
        tasks.append(_safe_send(ws, message))
    if tasks:
        await asyncio.gather(*tasks, return_exceptions=True)


async def _safe_send(ws: WebSocket, message: str):
    """Send text to one WebSocket, remove on failure."""
    try:
        await ws.send_text(message)
    except Exception:
        async with glock:
            connections.pop(ws, None)


async def _broadcast_members():
    """Broadcast JSON member list and count to all clients."""
    async with glock:
        names = list(connections.values())
        count = len(names)
    payload = json.dumps({"type": "members", "data": names})
    await _broadcast(f"MEMBERS::{payload}")
    await _broadcast(f"COUNT::{count}")


# ── WebSocket Endpoint ──────────────────────────────────────

@app.websocket("/wschat")
async def ws_endpoint(ws: WebSocket):
    await ws.accept()

    # ── 1) Receive username as first text frame ──
    try:
        raw = await asyncio.wait_for(ws.receive_text(), timeout=30)
    except Exception:
        await ws.close(1008)
        return

    username = (raw.strip() or f"匿名{id(ws) % 10000}")[:20]

    async with glock:
        connections[ws] = username

    join_msg = f"SYSTEM::{username} 加入聊天室"
    log_to_file(join_msg)
    asyncio.create_task(_broadcast(join_msg))
    asyncio.create_task(_broadcast_members())

    # ── 2) Message loop ──
    try:
        while True:
            text = await ws.receive_text()
            if not text or not text.strip():
                continue
            # Rate limit check
            if not await _check_msg_rate(ws):
                asyncio.create_task(_safe_send(ws, "SYSTEM::⚠️ 发送太频繁，请稍后再试"))
                continue
            msg = f"TEXT::{username}::{text.strip()}"
            log_to_file(msg)
            asyncio.create_task(_broadcast(msg))
    except (WebSocketDisconnect, Exception):
        pass
    finally:
        async with glock:
            name = connections.pop(ws, None) or "匿名用户"
        leave_msg = f"SYSTEM::{name} 离开聊天室"
        log_to_file(leave_msg)
        asyncio.create_task(_broadcast(leave_msg))
        asyncio.create_task(_broadcast_members())


# ── Nightly Cleanup ─────────────────────────────────────────

@app.post("/api/nightly-cleanup")
async def nightly_cleanup():
    """Backup chat history to backups/ dir, then clear the log.
    Also purge uploaded files older than 24h.
    Called by cron / systemd timer.
    """
    now = datetime.now()
    backup_name = f"chat_history_{now.strftime('%Y%m%d_%H%M%S')}.txt"
    backup_path = os.path.join(BACKUP_DIR, backup_name)

    if os.path.exists(CHAT_LOG_FILE) and os.path.getsize(CHAT_LOG_FILE) > 0:
        shutil.copy2(CHAT_LOG_FILE, backup_path)
        # Clear log
        open(CHAT_LOG_FILE, "w", encoding="utf-8").close()

    # Clean old uploads (> 24h)
    now_ts = time.time()
    cleaned = 0
    for fname in os.listdir(UPLOAD_DIR):
        fpath = os.path.join(UPLOAD_DIR, fname)
        if os.path.isfile(fpath) and now_ts - os.path.getmtime(fpath) > 86400:
            try:
                os.remove(fpath)
                cleaned += 1
            except Exception:
                pass

    return {"status": "ok", "backup": backup_name, "cleaned_uploads": cleaned}


# ── Main ────────────────────────────────────────────────────

if __name__ == "__main__":
    import uvicorn
    uvicorn.run("chat:app", host="0.0.0.0", port=8202, reload=False, log_level="info")