backend/app/api/email.py

import uuid
import logging
from typing import List, Optional
from fastapi import APIRouter, Depends, HTTPException, status, Query
from fastapi.responses import StreamingResponse
from sqlalchemy import select, update, desc
from sqlalchemy.ext.asyncio import AsyncSession

from app.database.postgres import get_db, Email, User
from app.database.mongo import mongo_db
from app.auth.security import get_current_user
from app.agent_logic.agent import stream_draft_reply
from app.tasks.email_tasks import send_email_task
from pydantic import BaseModel, EmailStr

logger = logging.getLogger(__name__)
router = APIRouter(prefix="/api/emails", tags=["emails"])

class EmailListItem(BaseModel):
    id: str
    subject: Optional[str]
    sender: str
    received_at: str
    category: str
    is_read: bool
    is_starred: bool
    thread_id: Optional[str]
    has_attachment: bool

class EmailDetailResponse(BaseModel):
    id: str
    subject: Optional[str]
    sender: str
    received_at: str
    category: str
    is_read: bool
    is_starred: bool
    thread_id: Optional[str]
    has_attachment: bool
    body_html: str
    body_text: str
    summary: Optional[str] = None
    action_items: List[str] = []

class SendEmailPayload(BaseModel):
    to: EmailStr
    subject: str
    body: str

@router.get("", response_model=List[EmailListItem])
async def list_emails(
    category: Optional[str] = Query(None, description="Filter by category (urgent, personal, invoice, newsletter)"),
    db: AsyncSession = Depends(get_db),
    current_user: User = Depends(get_current_user)
):
    """Retrieve email metadata for the current user."""
    query = select(Email).where(Email.user_id == current_user.id)
    if category:
        query = query.where(Email.category == category)
        
    query = query.order_by(desc(Email.received_at))
    result = await db.execute(query)
    emails = result.scalars().all()
    
    return [
        EmailListItem(
            id=str(e.id),
            subject=e.subject,
            sender=e.sender,
            received_at=e.received_at.isoformat(),
            category=e.category,
            is_read=e.is_read,
            is_starred=e.is_starred,
            thread_id=e.thread_id,
            has_attachment=e.has_attachment
        )
        for e in emails
    ]

@router.get("/{email_id}", response_model=EmailDetailResponse)
async def get_email_details(
    email_id: str,
    db: AsyncSession = Depends(get_db),
    current_user: User = Depends(get_current_user)
):
    """Retrieve complete email meta, body, summary, and action items."""
    try:
        email_uuid = uuid.UUID(email_id)
    except ValueError:
        raise HTTPException(status_code=400, detail="Invalid email UUID format")
        
    # Query PostgreSQL
    result = await db.execute(
        select(Email).where(Email.id == email_uuid, Email.user_id == current_user.id)
    )
    email = result.scalars().first()
    if not email:
        raise HTTPException(status_code=404, detail="Email not found")
        
    # Query MongoDB
    body_doc = await mongo_db.email_bodies.find_one({"email_id": email_id})
    if not body_doc:
        body_doc = {}
        
    # Mark as read since user opened details
    if not email.is_read:
        email.is_read = True
        db.add(email)
        await db.commit()

    return EmailDetailResponse(
        id=str(email.id),
        subject=email.subject,
        sender=email.sender,
        received_at=email.received_at.isoformat(),
        category=email.category,
        is_read=email.is_read,
        is_starred=email.is_starred,
        thread_id=email.thread_id,
        has_attachment=email.has_attachment,
        body_html=body_doc.get("body_html", ""),
        body_text=body_doc.get("body_text", ""),
        summary=body_doc.get("summary"),
        action_items=body_doc.get("action_items", [])
    )

@router.post("/{email_id}/toggle-read")
async def toggle_read(
    email_id: str,
    is_read: bool,
    db: AsyncSession = Depends(get_db),
    current_user: User = Depends(get_current_user)
):
    """Toggle read/unread status."""
    email_uuid = uuid.UUID(email_id)
    result = await db.execute(
        select(Email).where(Email.id == email_uuid, Email.user_id == current_user.id)
    )
    email = result.scalars().first()
    if not email:
        raise HTTPException(status_code=404, detail="Email not found")
        
    email.is_read = is_read
    db.add(email)
    await db.commit()
    return {"status": "success", "is_read": email.is_read}

@router.post("/{email_id}/toggle-star")
async def toggle_star(
    email_id: str,
    is_starred: bool,
    db: AsyncSession = Depends(get_db),
    current_user: User = Depends(get_current_user)
):
    """Toggle starred status."""
    email_uuid = uuid.UUID(email_id)
    result = await db.execute(
        select(Email).where(Email.id == email_uuid, Email.user_id == current_user.id)
    )
    email = result.scalars().first()
    if not email:
        raise HTTPException(status_code=404, detail="Email not found")
        
    email.is_starred = is_starred
    db.add(email)
    await db.commit()
    return {"status": "success", "is_starred": email.is_starred}

@router.post("/send")
async def send_email(
    payload: SendEmailPayload,
    current_user: User = Depends(get_current_user)
):
    """Enqueues send email background task."""
    task_payload = {
        "to": payload.to,
        "subject": payload.subject,
        "body": payload.body
    }
    task = send_email_task.delay(task_payload, str(current_user.id))
    return {"status": "queued", "task_id": task.id}

@router.get("/{email_id}/draft")
async def get_streaming_draft(
    email_id: str,
    tone: str = Query("professional", description="Tone of reply: professional, casual, short"),
    db: AsyncSession = Depends(get_db),
    current_user: User = Depends(get_current_user)
):
    """Server-Sent Events (SSE) endpoint to stream OpenAI generated draft response."""
    try:
        email_uuid = uuid.UUID(email_id)
    except ValueError:
        raise HTTPException(status_code=400, detail="Invalid email UUID format")
        
    # Verify email ownership
    result = await db.execute(
        select(Email).where(Email.id == email_uuid, Email.user_id == current_user.id)
    )
    email = result.scalars().first()
    if not email:
        raise HTTPException(status_code=404, detail="Email not found")
        
    body_doc = await mongo_db.email_bodies.find_one({"email_id": email_id})
    body_text = body_doc.get("body_text", "") if body_doc else ""

    async def sse_generator():
        try:
            async for chunk in stream_draft_reply(
                subject=email.subject,
                sender=email.sender,
                body=body_text,
                tone=tone,
                user_id=str(current_user.id)
            ):
                # Standard Server-Sent Events formatting
                yield f"data: {chunk}\n\n"
        except Exception as e:
            logger.error(f"SSE draft streaming failure: {e}")
            yield f"data: Error: {str(e)}\n\n"

    return StreamingResponse(sse_generator(), media_type="text/event-stream")