app/api/notifications.py

from fastapi import APIRouter, Depends, HTTPException
from sqlalchemy.orm import Session
from pydantic import BaseModel
from typing import Optional
from datetime import datetime, timezone

from app.db.base import get_db
from app.models.notification import Notification
from app.api.auth import get_current_user
from app.core.security import get_current_user_sse
from app.models.user import User, UserRole
from app.core.notifications import notification_manager
from fastapi.responses import StreamingResponse
import asyncio
import json

router = APIRouter(prefix="/notifications", tags=["Notifications"])


def _check_admin(user_id: int, db: Session):
    user = db.query(User).filter(User.id == user_id).first()
    if not user or user.role != UserRole.ADMIN:
        raise HTTPException(status_code=403, detail="Admin access required")
    return user


# ── Schemas ───────────────────────────────────────────────────
class NotificationCreate(BaseModel):
    type: str  # checkout_shipping | checkout_payment | checkout_complete
    title: str
    message: Optional[str] = None
    data: Optional[dict] = None


class NotificationOut(BaseModel):
    id: int
    type: str
    title: str
    message: Optional[str]
    data: Optional[dict]
    is_read: bool
    created_at: datetime

    class Config:
        from_attributes = True


# ── POST /notifications — create (no auth, called from checkout) ──
@router.post("/")
async def create_notification(payload: NotificationCreate, db: Session = Depends(get_db)):
    notif = Notification(
        type=payload.type,
        title=payload.title,
        message=payload.message,
        data=payload.data,
        is_read=False,
        created_at=datetime.now(timezone.utc),
    )
    db.add(notif)
    db.commit()
    db.refresh(notif)
    
    # Broadcast to active SSE clients
    await notification_manager.broadcast(
        NotificationOut.model_validate(notif).model_dump()
    )
    
    return {
        "isSuccess": True,
        "value": {"id": notif.id},
        "statusCode": 201,
    }


# ── GET /notifications — list for admin ──
@router.get("/")
async def list_notifications(
    skip: int = 0,
    limit: int = 50,
    unread_only: bool = False,
    current_user_id: int = Depends(get_current_user),
    db: Session = Depends(get_db),
):
    _check_admin(current_user_id, db)

    query = db.query(Notification)
    if unread_only:
        query = query.filter(Notification.is_read == False)  # noqa: E712
    query = query.order_by(Notification.created_at.desc())

    total = query.count()
    notifications = query.offset(skip).limit(limit).all()

    return {
        "isSuccess": True,
        "value": {
            "notifications": [
                NotificationOut.model_validate(n).model_dump() for n in notifications
            ],
            "total": total,
            "unread_count": db.query(Notification).filter(Notification.is_read == False).count(),  # noqa: E712
        },
        "statusCode": 200,
    }


# ── PUT /notifications/{id}/read — mark single as read ──
@router.put("/{notification_id}/read")
async def mark_read(
    notification_id: int,
    current_user_id: int = Depends(get_current_user),
    db: Session = Depends(get_db),
):
    _check_admin(current_user_id, db)

    notif = db.query(Notification).filter(Notification.id == notification_id).first()
    if not notif:
        raise HTTPException(status_code=404, detail="Notification not found")

    notif.is_read = True
    db.commit()
    return {"isSuccess": True, "value": None, "statusCode": 200}


# ── PUT /notifications/read-all — mark all as read ──
@router.put("/read-all")
async def mark_all_read(
    current_user_id: int = Depends(get_current_user),
    db: Session = Depends(get_db),
):
    _check_admin(current_user_id, db)

    db.query(Notification).filter(Notification.is_read == False).update(  # noqa: E712
        {"is_read": True}
    )
    db.commit()
    return {"isSuccess": True, "value": None, "statusCode": 200}


# ── GET /notifications/stream — SSE stream for admin ──
@router.get("/stream")
async def stream_notifications(
    current_user_id: int = Depends(get_current_user_sse),
    db: Session = Depends(get_db),
):
    """
    SSE endpoint for real-time notifications.
    Includes a heartbeat to keep the connection alive on Hugging Face / Vercel.
    """
    _check_admin(current_user_id, db)

    async def event_generator():
        queue = await notification_manager.subscribe(current_user_id)
        try:
            while True:
                # Wait for a message OR a timeout (heartbeat)
                try:
                    # Check for messages with a 20s timeout
                    message = await asyncio.wait_for(queue.get(), timeout=20.0)
                    yield f"data: {message}\n\n"
                except asyncio.TimeoutError:
                    # Send a heartbeat comment to keep the connection alive
                    yield ": heartbeat\n\n"
        finally:
            await notification_manager.unsubscribe(current_user_id, queue)

    return StreamingResponse(
        event_generator(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "X-Accel-Buffering": "no",  # Disable buffering for Nginx/Proxies
        },
    )