app/api/v1/webhook_socket.py

from __future__ import annotations

import asyncio
import hmac as hmac_mod
import json

from fastapi import (
    APIRouter,
    HTTPException,
    Request,
    WebSocket,
    WebSocketDisconnect,
)

from app.models.schemas import (
    ChannelCreateRequest,
    ChannelCreateResponse,
    ChannelDeleteResponse,
    ChannelInfoResponse,
    ChannelListItem,
    ChannelListResponse,
    WebhookResponse,
    WebhookSocketStatsResponse,
)
from app.services.webhook_socket_service import get_manager, verify_signature

router = APIRouter()
manager = get_manager()


@router.post("/channels", response_model=ChannelCreateResponse, summary="Create a webhook channel")
async def create_channel(
    body: ChannelCreateRequest,
    request: Request,
):
    if body.channel_id and manager.get_channel(body.channel_id):
        raise HTTPException(status_code=409, detail=f"Channel '{body.channel_id}' already exists")

    ch = manager.create_channel(
        channel_id=body.channel_id,
        secret=body.secret,
        buffer_size=body.buffer_size,
    )

    host = request.headers.get("host", "localhost:7860")
    scheme = request.headers.get("x-forwarded-proto", "http")
    ws_scheme = "wss" if scheme == "https" else "ws"

    return ChannelCreateResponse(
        channel_id=ch.channel_id,
        webhook_url=f"{scheme}://{host}/api/v1/webhook/{ch.channel_id}",
        ws_url=f"{ws_scheme}://{host}/api/v1/ws/{ch.channel_id}",
        secret=ch.secret,
        buffer_size=ch.buffer_size,
    )


@router.get("/channels", response_model=ChannelListResponse, summary="List all webhook channels")
async def list_channels():
    channels = []
    for cid, ch in manager.channels.items():
        channels.append(ChannelListItem(
            channel_id=cid,
            subscribers=len(ch.subscribers),
            messages=ch.message_count,
            buffered=len(ch.history),
            created_at=ch.created_at,
            last_activity=ch.last_activity,
        ))
    return ChannelListResponse(channels=channels)


@router.get("/channels/{channel_id}", response_model=ChannelInfoResponse, summary="Get channel info")
async def channel_info(
    channel_id: str,
):
    ch = manager.get_channel(channel_id)
    if not ch:
        raise HTTPException(status_code=404, detail="Channel not found")
    return ChannelInfoResponse(
        channel_id=ch.channel_id,
        subscribers=len(ch.subscribers),
        messages=ch.message_count,
        buffered=len(ch.history),
        buffer_size=ch.buffer_size,
        has_secret=ch.secret is not None,
        created_at=ch.created_at,
        last_activity=ch.last_activity,
    )


@router.delete("/channels/{channel_id}", response_model=ChannelDeleteResponse, summary="Delete a channel")
async def delete_channel(
    channel_id: str,
):
    if manager.delete_channel(channel_id):
        return ChannelDeleteResponse(deleted=channel_id)
    raise HTTPException(status_code=404, detail="Channel not found")


@router.post("/webhook/{channel_id}", response_model=WebhookResponse, summary="Send webhook payload to channel")
async def handle_webhook(
    channel_id: str,
    request: Request,
):
    ch = manager.get_channel(channel_id)
    if not ch:
        raise HTTPException(status_code=404, detail="Channel not found")

    raw_body = await request.body()

    if ch.secret:
        sig_header = request.headers.get("X-Signature-256") or request.headers.get("X-Hub-Signature-256", "")
        if not sig_header:
            raise HTTPException(status_code=401, detail="Missing signature header")
        if not verify_signature(ch.secret, raw_body, sig_header):
            raise HTTPException(status_code=401, detail="Invalid signature")

    content_type = request.headers.get("content-type", "")

    if "json" in content_type:
        try:
            payload = json.loads(raw_body)
        except json.JSONDecodeError:
            payload = raw_body.decode(errors="replace")
    elif "x-www-form-urlencoded" in content_type:
        form = await request.form()
        payload = dict(form)
    else:
        try:
            payload = json.loads(raw_body)
        except json.JSONDecodeError:
            payload = raw_body.decode(errors="replace")

    fwd_headers: dict[str, str] = {}
    for h in ("X-GitHub-Event", "X-GitHub-Delivery", "X-Event-Type",
              "X-Webhook-Name", "Content-Type", "User-Agent"):
        if h in request.headers:
            fwd_headers[h] = request.headers[h]

    sent = await manager.publish(channel_id, payload, fwd_headers)

    ch_ref = manager.get_channel(channel_id)
    return WebhookResponse(
        status="delivered",
        channel=channel_id,
        subscribers_notified=max(sent, 0),
        message_id=ch_ref.message_count if ch_ref else 0,
    )


@router.post("/hook/{channel_id}", response_model=WebhookResponse, summary="Send webhook payload (short alias)")
async def handle_webhook_short(
    channel_id: str,
    request: Request,
):
    return await handle_webhook(channel_id, request)


@router.get("/ws/{channel_id}")
async def websocket_endpoint(
    channel_id: str,
    websocket: WebSocket,
    secret: str = "",
):
    ch = manager.get_channel(channel_id)
    if not ch:
        await websocket.accept()
        await websocket.send_json({"event": "error", "message": "channel not found"})
        await websocket.close(code=4404)
        return

    if ch.secret:
        if not hmac_mod.compare_digest(secret, ch.secret):
            await websocket.accept()
            await websocket.send_json({"event": "error", "message": "unauthorized"})
            await websocket.close(code=4401)
            return

    await websocket.accept()

    q = manager.subscribe(channel_id, websocket)
    if q is None:
        await websocket.send_json({"event": "error", "message": "subscribe failed"})
        await websocket.close()
        return

    await websocket.send_json({
        "event": "connected",
        "channel": channel_id,
        "message": f"Listening on channel '{channel_id}'",
        "buffered": len(ch.history),
    })

    async def forward_to_ws():
        try:
            while True:
                try:
                    msg = await asyncio.wait_for(q.get(), timeout=30)
                except asyncio.TimeoutError:
                    try:
                        await websocket.send_json({"event": "ping"})
                    except Exception:
                        break
                    continue
                try:
                    await websocket.send_json(msg)
                except Exception:
                    break
        except asyncio.CancelledError:
            pass

    async def read_from_ws():
        try:
            while True:
                try:
                    data = await websocket.receive_text()
                    try:
                        msg_data = json.loads(data)
                        if msg_data.get("type") == "ping":
                            await websocket.send_json({"event": "pong"})
                    except json.JSONDecodeError:
                        pass
                except WebSocketDisconnect:
                    break
        except asyncio.CancelledError:
            pass

    fwd_task = asyncio.create_task(forward_to_ws())
    read_task = asyncio.create_task(read_from_ws())

    try:
        done, pending = await asyncio.wait(
            [fwd_task, read_task], return_when=asyncio.FIRST_COMPLETED
        )
    finally:
        fwd_task.cancel()
        read_task.cancel()
        manager.unsubscribe(channel_id, websocket)


@router.get("/webhook-socket/stats", response_model=WebhookSocketStatsResponse, summary="Webhook/socket server stats")
async def ws_stats():
    return manager.stats()