PROJECT-CONTORL/backend/app/api/attachments/router.py

import uuid
import logging
from datetime import datetime
from io import BytesIO
from fastapi import APIRouter, Depends, HTTPException, UploadFile, File, Request
from fastapi.responses import FileResponse, Response
from sqlalchemy.orm import Session
from typing import Optional

from app.core.database import get_db
from app.middleware.auth import get_current_user, check_task_access, check_task_edit_access
from app.models import User, Task, Project, Attachment, AttachmentVersion, EncryptionKey, AuditAction
from app.schemas.attachment import (
    AttachmentResponse, AttachmentListResponse, AttachmentDetailResponse,
    AttachmentVersionResponse, VersionHistoryResponse
)
from app.services.file_storage_service import file_storage_service
from app.services.audit_service import AuditService
from app.services.watermark_service import watermark_service
from app.services.encryption_service import (
    encryption_service,
    MasterKeyNotConfiguredError,
    DecryptionError,
)

logger = logging.getLogger(__name__)

router = APIRouter(prefix="/api", tags=["attachments"])


def get_task_with_access_check(db: Session, task_id: str, current_user: User, require_edit: bool = False) -> Task:
    """Get task and verify access permissions."""
    task = db.query(Task).filter(Task.id == task_id).first()
    if not task:
        raise HTTPException(status_code=404, detail="Task not found")

    # Get project for access check
    project = db.query(Project).filter(Project.id == task.project_id).first()
    if not project:
        raise HTTPException(status_code=404, detail="Project not found")

    # Check access permission
    if not check_task_access(current_user, task, project):
        raise HTTPException(status_code=403, detail="Access denied to this task")

    # Check edit permission if required
    if require_edit and not check_task_edit_access(current_user, task, project):
        raise HTTPException(status_code=403, detail="Edit access denied to this task")

    return task


def get_attachment_with_access_check(
    db: Session, attachment_id: str, current_user: User, require_edit: bool = False
) -> Attachment:
    """Get attachment and verify access permissions."""
    attachment = db.query(Attachment).filter(
        Attachment.id == attachment_id,
        Attachment.is_deleted == False
    ).first()
    if not attachment:
        raise HTTPException(status_code=404, detail="Attachment not found")

    # Get task and project for access check
    task = db.query(Task).filter(Task.id == attachment.task_id).first()
    if not task:
        raise HTTPException(status_code=404, detail="Task not found")

    project = db.query(Project).filter(Project.id == task.project_id).first()
    if not project:
        raise HTTPException(status_code=404, detail="Project not found")

    # Check access permission
    if not check_task_access(current_user, task, project):
        raise HTTPException(status_code=403, detail="Access denied to this attachment")

    # Check edit permission if required
    if require_edit and not check_task_edit_access(current_user, task, project):
        raise HTTPException(status_code=403, detail="Edit access denied to this attachment")

    return attachment


def attachment_to_response(attachment: Attachment) -> AttachmentResponse:
    """Convert Attachment model to response."""
    return AttachmentResponse(
        id=attachment.id,
        task_id=attachment.task_id,
        filename=attachment.filename,
        original_filename=attachment.original_filename,
        mime_type=attachment.mime_type,
        file_size=attachment.file_size,
        current_version=attachment.current_version,
        is_encrypted=attachment.is_encrypted,
        uploaded_by=attachment.uploaded_by,
        uploader_name=attachment.uploader.name if attachment.uploader else None,
        created_at=attachment.created_at,
        updated_at=attachment.updated_at
    )


def version_to_response(version: AttachmentVersion) -> AttachmentVersionResponse:
    """Convert AttachmentVersion model to response."""
    return AttachmentVersionResponse(
        id=version.id,
        version=version.version,
        file_size=version.file_size,
        checksum=version.checksum,
        uploaded_by=version.uploaded_by,
        uploader_name=version.uploader.name if version.uploader else None,
        created_at=version.created_at
    )


def should_encrypt_file(project: Project, db: Session) -> tuple[bool, Optional[EncryptionKey]]:
    """
    Determine if a file should be encrypted based on project security level.

    Returns:
        Tuple of (should_encrypt, encryption_key)
    """
    # Only encrypt for confidential projects
    if project.security_level != "confidential":
        return False, None

    # Check if encryption is available
    if not encryption_service.is_encryption_available():
        logger.warning(
            f"Project {project.id} is confidential but encryption is not configured. "
            "Files will be stored unencrypted."
        )
        return False, None

    # Get active encryption key
    active_key = db.query(EncryptionKey).filter(
        EncryptionKey.is_active == True
    ).first()

    if not active_key:
        logger.warning(
            f"Project {project.id} is confidential but no active encryption key exists. "
            "Files will be stored unencrypted. Create a key using /api/admin/encryption-keys/rotate"
        )
        return False, None

    return True, active_key


@router.post("/tasks/{task_id}/attachments", response_model=AttachmentResponse)
async def upload_attachment(
    task_id: str,
    request: Request,
    file: UploadFile = File(...),
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user)
):
    """
    Upload a file attachment to a task.

    For confidential projects, files are automatically encrypted using AES-256-GCM.
    """
    task = get_task_with_access_check(db, task_id, current_user, require_edit=True)

    # Get project to check security level
    project = db.query(Project).filter(Project.id == task.project_id).first()
    if not project:
        raise HTTPException(status_code=404, detail="Project not found")

    # Determine if encryption is needed
    should_encrypt, encryption_key = should_encrypt_file(project, db)

    # Check if attachment with same filename exists (for versioning)
    existing = db.query(Attachment).filter(
        Attachment.task_id == task_id,
        Attachment.original_filename == file.filename,
        Attachment.is_deleted == False
    ).first()

    if existing:
        # Create new version for existing attachment
        new_version = existing.current_version + 1

        if should_encrypt and encryption_key:
            # Read and encrypt file content
            file_content = await file.read()
            await file.seek(0)

            try:
                # Decrypt the encryption key
                raw_key = encryption_service.decrypt_key(encryption_key.key_data)
                # Encrypt the file
                encrypted_content = encryption_service.encrypt_bytes(file_content, raw_key)

                # Create a new UploadFile-like object with encrypted content
                encrypted_file = BytesIO(encrypted_content)
                encrypted_file.seek(0)

                # Save encrypted file using a modified approach
                file_path, _, checksum = await file_storage_service.save_file(
                    file=file,
                    project_id=task.project_id,
                    task_id=task_id,
                    attachment_id=existing.id,
                    version=new_version
                )

                # Overwrite with encrypted content
                with open(file_path, "wb") as f:
                    f.write(encrypted_content)

                file_size = len(encrypted_content)

                # Update existing attachment with encryption info
                existing.is_encrypted = True
                existing.encryption_key_id = encryption_key.id

                # Audit log for encryption
                AuditService.log_event(
                    db=db,
                    event_type="attachment.encrypt",
                    resource_type="attachment",
                    action=AuditAction.UPDATE,
                    user_id=current_user.id,
                    resource_id=existing.id,
                    changes=[
                        {"field": "is_encrypted", "old_value": False, "new_value": True},
                        {"field": "encryption_key_id", "old_value": None, "new_value": encryption_key.id},
                    ],
                    request_metadata=getattr(request.state, "audit_metadata", None),
                )

            except Exception as e:
                logger.error(f"Failed to encrypt file for attachment {existing.id}: {e}")
                raise HTTPException(
                    status_code=500,
                    detail="Failed to encrypt file. Please try again."
                )
        else:
            # Save file without encryption
            file_path, file_size, checksum = await file_storage_service.save_file(
                file=file,
                project_id=task.project_id,
                task_id=task_id,
                attachment_id=existing.id,
                version=new_version
            )

        # Create version record
        version = AttachmentVersion(
            id=str(uuid.uuid4()),
            attachment_id=existing.id,
            version=new_version,
            file_path=file_path,
            file_size=file_size,
            checksum=checksum,
            uploaded_by=current_user.id
        )
        db.add(version)

        # Update attachment
        existing.current_version = new_version
        existing.file_size = file_size
        existing.updated_at = version.created_at

        # Audit log
        AuditService.log_event(
            db=db,
            event_type="attachment.upload",
            resource_type="attachment",
            action=AuditAction.UPDATE,
            user_id=current_user.id,
            resource_id=existing.id,
            changes=[{"field": "version", "old_value": new_version - 1, "new_value": new_version}],
            request_metadata=getattr(request.state, "audit_metadata", None)
        )

        db.commit()
        db.refresh(existing)

        return attachment_to_response(existing)

    # Create new attachment
    attachment_id = str(uuid.uuid4())
    is_encrypted = False
    encryption_key_id = None

    if should_encrypt and encryption_key:
        # Read and encrypt file content
        file_content = await file.read()
        await file.seek(0)

        try:
            # Decrypt the encryption key
            raw_key = encryption_service.decrypt_key(encryption_key.key_data)
            # Encrypt the file
            encrypted_content = encryption_service.encrypt_bytes(file_content, raw_key)

            # Save file first to get path
            file_path, _, checksum = await file_storage_service.save_file(
                file=file,
                project_id=task.project_id,
                task_id=task_id,
                attachment_id=attachment_id,
                version=1
            )

            # Overwrite with encrypted content
            with open(file_path, "wb") as f:
                f.write(encrypted_content)

            file_size = len(encrypted_content)
            is_encrypted = True
            encryption_key_id = encryption_key.id

        except Exception as e:
            logger.error(f"Failed to encrypt new file: {e}")
            raise HTTPException(
                status_code=500,
                detail="Failed to encrypt file. Please try again."
            )
    else:
        # Save file without encryption
        file_path, file_size, checksum = await file_storage_service.save_file(
            file=file,
            project_id=task.project_id,
            task_id=task_id,
            attachment_id=attachment_id,
            version=1
        )

    # Get mime type from file storage validation
    extension = file_storage_service.get_extension(file.filename or "")
    mime_type = file.content_type or "application/octet-stream"

    # Create attachment record
    attachment = Attachment(
        id=attachment_id,
        task_id=task_id,
        filename=file.filename or "unnamed",
        original_filename=file.filename or "unnamed",
        mime_type=mime_type,
        file_size=file_size,
        current_version=1,
        is_encrypted=is_encrypted,
        encryption_key_id=encryption_key_id,
        uploaded_by=current_user.id
    )
    db.add(attachment)

    # Create version record
    version = AttachmentVersion(
        id=str(uuid.uuid4()),
        attachment_id=attachment_id,
        version=1,
        file_path=file_path,
        file_size=file_size,
        checksum=checksum,
        uploaded_by=current_user.id
    )
    db.add(version)

    # Audit log
    changes = [{"field": "filename", "old_value": None, "new_value": attachment.filename}]
    if is_encrypted:
        changes.append({"field": "is_encrypted", "old_value": None, "new_value": True})

    AuditService.log_event(
        db=db,
        event_type="attachment.upload",
        resource_type="attachment",
        action=AuditAction.CREATE,
        user_id=current_user.id,
        resource_id=attachment.id,
        changes=changes,
        request_metadata=getattr(request.state, "audit_metadata", None)
    )

    db.commit()
    db.refresh(attachment)

    return attachment_to_response(attachment)


@router.get("/tasks/{task_id}/attachments", response_model=AttachmentListResponse)
async def list_task_attachments(
    task_id: str,
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user)
):
    """List all attachments for a task."""
    task = get_task_with_access_check(db, task_id, current_user, require_edit=False)

    attachments = db.query(Attachment).filter(
        Attachment.task_id == task_id,
        Attachment.is_deleted == False
    ).order_by(Attachment.created_at.desc()).all()

    return AttachmentListResponse(
        attachments=[attachment_to_response(a) for a in attachments],
        total=len(attachments)
    )


@router.get("/attachments/{attachment_id}", response_model=AttachmentDetailResponse)
async def get_attachment(
    attachment_id: str,
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user)
):
    """Get attachment details with version history."""
    attachment = get_attachment_with_access_check(db, attachment_id, current_user, require_edit=False)

    versions = db.query(AttachmentVersion).filter(
        AttachmentVersion.attachment_id == attachment_id
    ).order_by(AttachmentVersion.version.desc()).all()

    return AttachmentDetailResponse(
        id=attachment.id,
        task_id=attachment.task_id,
        filename=attachment.filename,
        original_filename=attachment.original_filename,
        mime_type=attachment.mime_type,
        file_size=attachment.file_size,
        current_version=attachment.current_version,
        is_encrypted=attachment.is_encrypted,
        uploaded_by=attachment.uploaded_by,
        uploader_name=attachment.uploader.name if attachment.uploader else None,
        created_at=attachment.created_at,
        updated_at=attachment.updated_at,
        versions=[version_to_response(v) for v in versions]
    )


@router.get("/attachments/{attachment_id}/download")
async def download_attachment(
    attachment_id: str,
    version: Optional[int] = None,
    request: Request = None,
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user)
):
    """
    Download an attachment file with dynamic watermark.

    For encrypted files, the file is automatically decrypted before returning.
    """
    attachment = get_attachment_with_access_check(db, attachment_id, current_user, require_edit=False)

    # Get version to download
    target_version = version or attachment.current_version

    version_record = db.query(AttachmentVersion).filter(
        AttachmentVersion.attachment_id == attachment_id,
        AttachmentVersion.version == target_version
    ).first()

    if not version_record:
        raise HTTPException(status_code=404, detail=f"Version {target_version} not found")

    # Get file path
    file_path = file_storage_service.get_file_by_path(version_record.file_path)
    if not file_path:
        raise HTTPException(status_code=404, detail="File not found on disk")

    # Audit log
    download_time = datetime.now()
    AuditService.log_event(
        db=db,
        event_type="attachment.download",
        resource_type="attachment",
        action=AuditAction.UPDATE,  # Using UPDATE as there's no DOWNLOAD action
        user_id=current_user.id,
        resource_id=attachment.id,
        changes=[{"field": "downloaded_version", "old_value": None, "new_value": target_version}],
        request_metadata=getattr(request.state, "audit_metadata", None) if request else None
    )
    db.commit()

    # Read file content
    with open(file_path, "rb") as f:
        file_bytes = f.read()

    # Decrypt if encrypted
    if attachment.is_encrypted:
        if not attachment.encryption_key_id:
            raise HTTPException(
                status_code=500,
                detail="Encrypted file is missing encryption key reference"
            )

        encryption_key = db.query(EncryptionKey).filter(
            EncryptionKey.id == attachment.encryption_key_id
        ).first()

        if not encryption_key:
            raise HTTPException(
                status_code=500,
                detail="Encryption key not found for this file"
            )

        try:
            # Decrypt the encryption key
            raw_key = encryption_service.decrypt_key(encryption_key.key_data)
            # Decrypt the file
            file_bytes = encryption_service.decrypt_bytes(file_bytes, raw_key)

            # Audit log for decryption
            AuditService.log_event(
                db=db,
                event_type="attachment.decrypt",
                resource_type="attachment",
                action=AuditAction.UPDATE,
                user_id=current_user.id,
                resource_id=attachment.id,
                changes=[{"field": "decrypted_for_download", "old_value": None, "new_value": True}],
                request_metadata=getattr(request.state, "audit_metadata", None) if request else None,
            )
            db.commit()

        except DecryptionError as e:
            logger.error(f"Failed to decrypt attachment {attachment_id}: {e}")
            raise HTTPException(
                status_code=500,
                detail="Failed to decrypt file. The file may be corrupted."
            )
        except MasterKeyNotConfiguredError:
            raise HTTPException(
                status_code=500,
                detail="Encryption is not configured. Cannot decrypt file."
            )
        except Exception as e:
            logger.error(f"Unexpected error decrypting attachment {attachment_id}: {e}")
            raise HTTPException(
                status_code=500,
                detail="Failed to decrypt file."
            )

    # Check if watermark should be applied
    mime_type = attachment.mime_type or ""
    if watermark_service.supports_watermark(mime_type):
        try:
            # Apply watermark based on file type
            if watermark_service.is_supported_image(mime_type):
                watermarked_bytes, output_format = watermark_service.add_image_watermark(
                    image_bytes=file_bytes,
                    user_name=current_user.name,
                    employee_id=current_user.employee_id,
                    download_time=download_time
                )
                # Update mime type based on output format
                output_mime_type = f"image/{output_format}"
                # Update filename extension if format changed
                original_filename = attachment.original_filename
                if output_format == "png" and not original_filename.lower().endswith(".png"):
                    original_filename = original_filename.rsplit(".", 1)[0] + ".png"

                return Response(
                    content=watermarked_bytes,
                    media_type=output_mime_type,
                    headers={
                        "Content-Disposition": f'attachment; filename="{original_filename}"'
                    }
                )

            elif watermark_service.is_supported_pdf(mime_type):
                watermarked_bytes = watermark_service.add_pdf_watermark(
                    pdf_bytes=file_bytes,
                    user_name=current_user.name,
                    employee_id=current_user.employee_id,
                    download_time=download_time
                )

                return Response(
                    content=watermarked_bytes,
                    media_type="application/pdf",
                    headers={
                        "Content-Disposition": f'attachment; filename="{attachment.original_filename}"'
                    }
                )

        except Exception as e:
            # If watermarking fails, log the error but still return the file
            logger.warning(
                f"Watermarking failed for attachment {attachment_id}: {str(e)}. "
                "Returning file without watermark."
            )

    # Return file (decrypted if needed, without watermark for unsupported types)
    return Response(
        content=file_bytes,
        media_type=attachment.mime_type,
        headers={
            "Content-Disposition": f'attachment; filename="{attachment.original_filename}"'
        }
    )


@router.delete("/attachments/{attachment_id}")
async def delete_attachment(
    attachment_id: str,
    request: Request,
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user)
):
    """Soft delete an attachment."""
    attachment = get_attachment_with_access_check(db, attachment_id, current_user, require_edit=True)

    # Soft delete
    attachment.is_deleted = True

    # Audit log
    AuditService.log_event(
        db=db,
        event_type="attachment.delete",
        resource_type="attachment",
        action=AuditAction.DELETE,
        user_id=current_user.id,
        resource_id=attachment.id,
        changes=[{"field": "is_deleted", "old_value": False, "new_value": True}],
        request_metadata=getattr(request.state, "audit_metadata", None)
    )

    db.commit()

    return {"detail": "Attachment deleted", "id": attachment_id}


@router.get("/attachments/{attachment_id}/versions", response_model=VersionHistoryResponse)
async def get_version_history(
    attachment_id: str,
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user)
):
    """Get version history for an attachment."""
    attachment = get_attachment_with_access_check(db, attachment_id, current_user, require_edit=False)

    versions = db.query(AttachmentVersion).filter(
        AttachmentVersion.attachment_id == attachment_id
    ).order_by(AttachmentVersion.version.desc()).all()

    return VersionHistoryResponse(
        attachment_id=attachment.id,
        filename=attachment.filename,
        versions=[version_to_response(v) for v in versions],
        total=len(versions)
    )


@router.post("/attachments/{attachment_id}/restore/{version}")
async def restore_version(
    attachment_id: str,
    version: int,
    request: Request,
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user)
):
    """Restore an attachment to a specific version."""
    attachment = get_attachment_with_access_check(db, attachment_id, current_user, require_edit=True)

    version_record = db.query(AttachmentVersion).filter(
        AttachmentVersion.attachment_id == attachment_id,
        AttachmentVersion.version == version
    ).first()

    if not version_record:
        raise HTTPException(status_code=404, detail=f"Version {version} not found")

    old_version = attachment.current_version
    attachment.current_version = version
    attachment.file_size = version_record.file_size

    # Audit log
    AuditService.log_event(
        db=db,
        event_type="attachment.restore",
        resource_type="attachment",
        action=AuditAction.RESTORE,
        user_id=current_user.id,
        resource_id=attachment.id,
        changes=[{"field": "current_version", "old_value": old_version, "new_value": version}],
        request_metadata=getattr(request.state, "audit_metadata", None)
    )

    db.commit()

    return {"detail": f"Restored to version {version}", "current_version": version}