api/reports.py

import logging
from fastapi import APIRouter, Depends, HTTPException, BackgroundTasks
from sqlalchemy.orm import Session
from typing import List, Optional
from datetime import datetime

from api.auth import get_current_user
from models import db_models
from models.schemas import ReportGenerateRequest, ReportResponse, ReportFormatSuggestionResponse
from core.database import get_db, SessionLocal
from api.websocket_routes import manager
from services.report_service import report_service
from core import constants

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

async def run_report_generation(report_id: int, request: ReportGenerateRequest, user_id: int):
    """Background task for report generation"""
    db = SessionLocal()
    connection_id = f"user_{user_id}"
    try:
        db_report = db.query(db_models.Report).filter(db_models.Report.id == report_id).first()
        if not db_report: return

        # Call AI service
        content = await report_service.generate_report(
            file_key=request.file_key,
            text_input=request.text_input,
            format_key=request.format_key,
            custom_prompt=request.custom_prompt,
            language=request.language
        )

        if not content:
            raise Exception("AI failed to generate report content")

        if not db_report.title or "Report-" not in db_report.title:
            # Extract title (usually the first line)
            extracted_title = content.split('\n')[0].replace('#', '').strip()
            if not extracted_title or len(extracted_title) < 3:
                extracted_title = f"Report {request.format_key}"
            db_report.title = extracted_title
        db_report.content = content
        db_report.status = "completed"
        db.commit()

        # Notify via WebSocket
        await manager.send_result(connection_id, {
            "type": "report",
            "id": db_report.id,
            "status": "completed",
            "title": db_report.title
        })

    except Exception as e:
        logger.error(f"Background report generation failed: {e}")
        db_report = db.query(db_models.Report).filter(db_models.Report.id == report_id).first()
        if db_report:
            db_report.status = "failed"
            db_report.error_message = str(e)
            db.commit()
        await manager.send_error(connection_id, f"Report generation failed: {str(e)}")
    finally:
        db.close()

@router.get("/config")
async def get_report_config():
    """Returns available formats and languages for report generation."""
    return {
        "formats": constants.REPORT_FORMAT_OPTIONS,
        "languages": constants.LANGUAGES
    }

@router.get("/suggest-formats", response_model=ReportFormatSuggestionResponse)
async def suggest_formats(
    file_key: Optional[str] = None,
    text_input: Optional[str] = None,
    language: str = "Japanese",
    current_user: db_models.User = Depends(get_current_user)):
    """
    Get 4 AI-suggested report formats based on content.
    """
    suggestions = await report_service.generate_format_suggestions(
        file_key=file_key,
        text_input=text_input,
        language=language
    )
    return {"suggestions": suggestions}

@router.post("/generate", response_model=ReportResponse)
async def generate_report(
    request: ReportGenerateRequest,
    background_tasks: BackgroundTasks,
    current_user: db_models.User = Depends(get_current_user),
    db: Session = Depends(get_db)
):
    """
    Initiates report generation in the background.
    """
    source_id = None
    if request.file_key:
        source = db.query(db_models.Source).filter(
            db_models.Source.s3_key == request.file_key,
            db_models.Source.user_id == current_user.id
        ).first()
        if not source:
            raise HTTPException(status_code=403, detail="Not authorized to access this file")
        source_id = source.id

    # Create initial processing record
    file_base = request.file_key.split('/')[-1].rsplit('.', 1)[0] if request.file_key else None
    title = f"Report-{file_base}" if file_base else f"Report {request.format_key} {datetime.utcnow().strftime('%Y-%m-%d %H:%M')}"
    db_report = db_models.Report(
        title=title,
        format_key=request.format_key,
        user_id=current_user.id,
        source_id=source_id,
        status="processing"
    )
    db.add(db_report)
    db.commit()
    db.refresh(db_report)

    # Offload to background task
    background_tasks.add_task(run_report_generation, db_report.id, request, current_user.id)

    return db_report

@router.get("/list", response_model=List[ReportResponse])
async def list_reports(
    current_user: db_models.User = Depends(get_current_user),
    db: Session = Depends(get_db)
):
    """
    Lists all reports for the current user.
    """
    try:
        reports = db.query(db_models.Report).filter(
            db_models.Report.user_id == current_user.id
        ).order_by(db_models.Report.created_at.desc()).all()
        return [ReportResponse.model_validate(r) for r in reports]
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@router.get("/{report_id}", response_model=ReportResponse)
async def get_report(
    report_id: int,
    current_user: db_models.User = Depends(get_current_user),
    db: Session = Depends(get_db)
):
    """
    Retrieves a specific report.
    """
    report = db.query(db_models.Report).filter(
        db_models.Report.id == report_id,
        db_models.Report.user_id == current_user.id
    ).first()
    
    if not report:
        raise HTTPException(status_code=404, detail="Report not found")
    
    return ReportResponse.model_validate(report)

@router.delete("/{report_id}")
async def delete_report(
    report_id: int,
    current_user: db_models.User = Depends(get_current_user),
    db: Session = Depends(get_db)
):
    """
    Deletes a specific report.
    """
    report = db.query(db_models.Report).filter(
        db_models.Report.id == report_id,
        db_models.Report.user_id == current_user.id
    ).first()
    
    if not report:
        raise HTTPException(status_code=404, detail="Report not found")
    
    db.delete(report)
    db.commit()
    return {"message": "Report deleted successfully"}