src/api/services/report_service.py

"""
Service for working with intelligence reports.
"""
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.future import select
from sqlalchemy import update, delete, func, desc, and_, or_
from typing import List, Optional, Dict, Any, Union
import logging
from datetime import datetime

from src.models.report import Report, ReportType, ReportStatus
from src.models.threat import ThreatSeverity
from src.api.schemas import PaginationParams

# Configure logger
logger = logging.getLogger(__name__)


async def create_report(
    db: AsyncSession,
    title: str,
    summary: str,
    content: str,
    report_type: ReportType,
    report_id: str,
    status: ReportStatus = ReportStatus.DRAFT,
    severity: Optional[ThreatSeverity] = None,
    publish_date: Optional[datetime] = None,
    time_period_start: Optional[datetime] = None,
    time_period_end: Optional[datetime] = None,
    keywords: Optional[List[str]] = None,
    source_data: Optional[Dict[str, Any]] = None,
    author_id: Optional[int] = None,
) -> Report:
    """
    Create a new intelligence report.
    
    Args:
        db: Database session
        title: Report title
        summary: Report summary
        content: Report content
        report_type: Type of report
        report_id: Custom ID for the report
        status: Report status
        severity: Report severity
        publish_date: Publication date
        time_period_start: Start of time period covered
        time_period_end: End of time period covered
        keywords: List of keywords related to the report
        source_data: Sources and references
        author_id: ID of the report author
        
    Returns:
        Report: Created report
    """
    report = Report(
        title=title,
        summary=summary,
        content=content,
        report_type=report_type,
        report_id=report_id,
        status=status,
        severity=severity,
        publish_date=publish_date,
        time_period_start=time_period_start,
        time_period_end=time_period_end,
        keywords=keywords or [],
        source_data=source_data or {},
        author_id=author_id,
    )
    
    db.add(report)
    await db.commit()
    await db.refresh(report)
    
    return report


async def get_report_by_id(db: AsyncSession, report_id: int) -> Optional[Report]:
    """
    Get report by ID.
    
    Args:
        db: Database session
        report_id: Report ID
        
    Returns:
        Optional[Report]: Found report or None
    """
    result = await db.execute(
        select(Report).where(Report.id == report_id)
    )
    return result.scalars().first()


async def get_report_by_custom_id(db: AsyncSession, custom_id: str) -> Optional[Report]:
    """
    Get report by custom ID.
    
    Args:
        db: Database session
        custom_id: Custom report ID (e.g., RPT-2023-0001)
        
    Returns:
        Optional[Report]: Found report or None
    """
    result = await db.execute(
        select(Report).where(Report.report_id == custom_id)
    )
    return result.scalars().first()


async def get_reports(
    db: AsyncSession,
    pagination: PaginationParams,
    report_type: Optional[List[ReportType]] = None,
    status: Optional[List[ReportStatus]] = None,
    severity: Optional[List[ThreatSeverity]] = None,
    search_query: Optional[str] = None,
    keywords: Optional[List[str]] = None,
    author_id: Optional[int] = None,
    from_date: Optional[datetime] = None,
    to_date: Optional[datetime] = None,
) -> List[Report]:
    """
    Get reports with filtering.
    
    Args:
        db: Database session
        pagination: Pagination parameters
        report_type: Filter by report type
        status: Filter by status
        severity: Filter by severity
        search_query: Search in title and summary
        keywords: Filter by keywords
        author_id: Filter by author ID
        from_date: Filter by created_at >= from_date
        to_date: Filter by created_at <= to_date
        
    Returns:
        List[Report]: List of reports
    """
    query = select(Report)
    
    # Apply filters
    if report_type:
        query = query.where(Report.report_type.in_(report_type))
    
    if status:
        query = query.where(Report.status.in_(status))
    
    if severity:
        query = query.where(Report.severity.in_(severity))
    
    if search_query:
        search_filter = or_(
            Report.title.ilike(f"%{search_query}%"),
            Report.summary.ilike(f"%{search_query}%"),
            Report.content.ilike(f"%{search_query}%"),
        )
        query = query.where(search_filter)
    
    if keywords:
        # For JSON arrays, need to use a more complex query
        for keyword in keywords:
            query = query.where(Report.keywords.contains([keyword]))
    
    if author_id:
        query = query.where(Report.author_id == author_id)
    
    if from_date:
        query = query.where(Report.created_at >= from_date)
    
    if to_date:
        query = query.where(Report.created_at <= to_date)
    
    # Apply pagination
    query = query.order_by(desc(Report.created_at))
    query = query.offset((pagination.page - 1) * pagination.size).limit(pagination.size)
    
    result = await db.execute(query)
    return result.scalars().all()


async def count_reports(
    db: AsyncSession,
    report_type: Optional[List[ReportType]] = None,
    status: Optional[List[ReportStatus]] = None,
    severity: Optional[List[ThreatSeverity]] = None,
    search_query: Optional[str] = None,
    keywords: Optional[List[str]] = None,
    author_id: Optional[int] = None,
    from_date: Optional[datetime] = None,
    to_date: Optional[datetime] = None,
) -> int:
    """
    Count reports with filtering.
    
    Args are the same as get_reports, except pagination.
    
    Returns:
        int: Count of matching reports
    """
    query = select(func.count(Report.id))
    
    # Apply filters
    if report_type:
        query = query.where(Report.report_type.in_(report_type))
    
    if status:
        query = query.where(Report.status.in_(status))
    
    if severity:
        query = query.where(Report.severity.in_(severity))
    
    if search_query:
        search_filter = or_(
            Report.title.ilike(f"%{search_query}%"),
            Report.summary.ilike(f"%{search_query}%"),
            Report.content.ilike(f"%{search_query}%"),
        )
        query = query.where(search_filter)
    
    if keywords:
        # For JSON arrays, need to use a more complex query
        for keyword in keywords:
            query = query.where(Report.keywords.contains([keyword]))
    
    if author_id:
        query = query.where(Report.author_id == author_id)
    
    if from_date:
        query = query.where(Report.created_at >= from_date)
    
    if to_date:
        query = query.where(Report.created_at <= to_date)
    
    result = await db.execute(query)
    return result.scalar()


async def update_report(
    db: AsyncSession,
    report_id: int,
    title: Optional[str] = None,
    summary: Optional[str] = None,
    content: Optional[str] = None,
    report_type: Optional[ReportType] = None,
    status: Optional[ReportStatus] = None,
    severity: Optional[ThreatSeverity] = None,
    publish_date: Optional[datetime] = None,
    time_period_start: Optional[datetime] = None,
    time_period_end: Optional[datetime] = None,
    keywords: Optional[List[str]] = None,
    source_data: Optional[Dict[str, Any]] = None,
) -> Optional[Report]:
    """
    Update a report.
    
    Args:
        db: Database session
        report_id: Report ID
        Other args: Fields to update
        
    Returns:
        Optional[Report]: Updated report or None
    """
    report = await get_report_by_id(db, report_id)
    
    if not report:
        return None
    
    # Update fields if provided
    if title is not None:
        report.title = title
    
    if summary is not None:
        report.summary = summary
    
    if content is not None:
        report.content = content
    
    if report_type is not None:
        report.report_type = report_type
    
    if status is not None:
        report.status = status
    
    if severity is not None:
        report.severity = severity
    
    if publish_date is not None:
        report.publish_date = publish_date
    
    if time_period_start is not None:
        report.time_period_start = time_period_start
    
    if time_period_end is not None:
        report.time_period_end = time_period_end
    
    if keywords is not None:
        report.keywords = keywords
    
    if source_data is not None:
        report.source_data = source_data
    
    await db.commit()
    await db.refresh(report)
    
    return report


async def add_threat_to_report(
    db: AsyncSession,
    report_id: int,
    threat_id: int,
) -> Optional[Report]:
    """
    Add a threat to a report.
    
    Args:
        db: Database session
        report_id: Report ID
        threat_id: Threat ID
        
    Returns:
        Optional[Report]: Updated report or None
    """
    from src.api.services.threat_service import get_threat_by_id
    
    # Get report and threat
    report = await get_report_by_id(db, report_id)
    threat = await get_threat_by_id(db, threat_id)
    
    if not report or not threat:
        return None
    
    # Add threat to report
    report.threats.append(threat)
    await db.commit()
    await db.refresh(report)
    
    return report


async def remove_threat_from_report(
    db: AsyncSession,
    report_id: int,
    threat_id: int,
) -> Optional[Report]:
    """
    Remove a threat from a report.
    
    Args:
        db: Database session
        report_id: Report ID
        threat_id: Threat ID
        
    Returns:
        Optional[Report]: Updated report or None
    """
    from src.api.services.threat_service import get_threat_by_id
    
    # Get report and threat
    report = await get_report_by_id(db, report_id)
    threat = await get_threat_by_id(db, threat_id)
    
    if not report or not threat:
        return None
    
    # Remove threat from report
    if threat in report.threats:
        report.threats.remove(threat)
        await db.commit()
        await db.refresh(report)
    
    return report


async def publish_report(
    db: AsyncSession,
    report_id: int,
) -> Optional[Report]:
    """
    Publish a report.
    
    Args:
        db: Database session
        report_id: Report ID
        
    Returns:
        Optional[Report]: Updated report or None
    """
    report = await get_report_by_id(db, report_id)
    
    if not report:
        return None
    
    # Update status and publish date
    report.status = ReportStatus.PUBLISHED
    
    if not report.publish_date:
        report.publish_date = datetime.utcnow()
    
    await db.commit()
    await db.refresh(report)
    
    return report


async def archive_report(
    db: AsyncSession,
    report_id: int,
) -> Optional[Report]:
    """
    Archive a report.
    
    Args:
        db: Database session
        report_id: Report ID
        
    Returns:
        Optional[Report]: Updated report or None
    """
    report = await get_report_by_id(db, report_id)
    
    if not report:
        return None
    
    # Update status
    report.status = ReportStatus.ARCHIVED
    await db.commit()
    await db.refresh(report)
    
    return report