from datetime import datetime, timezone
from typing import List, Optional
from uuid import UUID

from common.database import DBConfig, Db
from common.logging import get_logger, span
from common.models.upload_session import (
    UploadSession,
    UploadSessionCreate,
    UploadSessionStatus,
)

logger = get_logger(__name__)


class UploadSessionDAO:
    """Data Access Object for upload_sessions in PostgreSQL."""

    def __init__(self, db: Optional[Db] = None):
        self.db = db if db else Db(DBConfig())
        logger.debug(
            "UploadSessionDAO initialized", extra={"component": "UploadSessionDAO"}
        )

    async def create_upload_session(
        self, session: UploadSessionCreate
    ) -> UploadSession:
        with span(
            logger,
            "create_upload_session",
            {
                "session_id": str(session.id),
                "user_id": session.user_id,
                "status": session.status.value,
            },
        ):
            columns = [
                "id",
                "user_id",
                "filename",
                "mime_type",
                "total_size_bytes",
                "uploaded_size_bytes",
                "status",
                "s3_multipart_upload_id",
                "s3_object_key",
                "failure_reason",
                "expires_at",
                "completed_at",
            ]
            values = (
                str(session.id),
                session.user_id,
                session.filename,
                session.mime_type,
                session.total_size_bytes,
                session.uploaded_size_bytes,
                session.status.value,
                session.s3_multipart_upload_id,
                session.s3_object_key,
                session.failure_reason,
                session.expires_at,
                session.completed_at,
            )

            created_id = await self.db.insert("upload_sessions", columns, values)
            if not created_id:
                raise RuntimeError("Failed to create upload session")

            created = await self.get_upload_session(session.id)
            if not created:
                raise RuntimeError("Upload session created but could not be retrieved")
            return created

    async def get_upload_session(
        self, session_id: UUID | str, user_id: Optional[int] = None
    ) -> Optional[UploadSession]:
        with span(
            logger,
            "get_upload_session",
            {"session_id": str(session_id), "user_id_filter": user_id},
        ):
            query = """
                SELECT id,
                       user_id,
                       filename,
                       mime_type,
                       total_size_bytes,
                       uploaded_size_bytes,
                       status,
                       s3_multipart_upload_id,
                       s3_object_key,
                       failure_reason,
                       expires_at,
                       completed_at,
                       created_at,
                       updated_at
                FROM upload_sessions
                WHERE id = %s
                """
            params: List[object] = [str(session_id)]
            if user_id is not None:
                query += " AND user_id = %s"
                params.append(user_id)

            row = await self.db.fetch_one(query, tuple(params))
            if not row:
                return None
            return self._row_to_upload_session(row)

    async def get_upload_session_by_object_key(
        self, object_key: str
    ) -> Optional[UploadSession]:
        with span(
            logger,
            "get_upload_session_by_object_key",
            {"s3_object_key": object_key},
        ):
            row = await self.db.fetch_one(
                """
                SELECT id,
                       user_id,
                       filename,
                       mime_type,
                       total_size_bytes,
                       uploaded_size_bytes,
                       status,
                       s3_multipart_upload_id,
                       s3_object_key,
                       failure_reason,
                       expires_at,
                       completed_at,
                       created_at,
                       updated_at
                FROM upload_sessions
                WHERE s3_object_key = %s
                """,
                (object_key,),
            )
            if not row:
                return None
            return self._row_to_upload_session(row)

    async def list_active_by_user(
        self, user_id: int, limit: int = 20, offset: int = 0
    ) -> List[UploadSession]:
        with span(
            logger,
            "list_active_by_user",
            {"user_id": user_id, "limit": limit, "offset": offset},
        ):
            rows = await self.db.fetch_all(
                """
                SELECT id,
                       user_id,
                       filename,
                       mime_type,
                       total_size_bytes,
                       uploaded_size_bytes,
                       status,
                       s3_multipart_upload_id,
                       s3_object_key,
                       failure_reason,
                       expires_at,
                       completed_at,
                       created_at,
                       updated_at
                FROM upload_sessions
                WHERE user_id = %s
                  AND status IN ('initiated', 'uploading', 'processing')
                ORDER BY created_at DESC
                LIMIT %s
                OFFSET %s
                """,
                (user_id, limit, offset),
            )
            if not rows:
                return []
            return [self._row_to_upload_session(row) for row in rows]

    async def set_progress(
        self,
        session_id: UUID | str,
        uploaded_size_bytes: int,
        status: Optional[UploadSessionStatus] = None,
    ) -> Optional[UploadSession]:
        with span(
            logger,
            "set_upload_session_progress",
            {
                "session_id": str(session_id),
                "uploaded_size_bytes": uploaded_size_bytes,
                "status": status.value if status else None,
            },
        ):
            set_columns = ["uploaded_size_bytes", "updated_at = CURRENT_TIMESTAMP"]
            set_values: List[object] = [uploaded_size_bytes]

            if status is not None:
                set_columns.append("status")
                set_values.append(status.value)

            await self.db.update(
                table="upload_sessions",
                set_columns=set_columns,
                set_values=set_values,
                where_clause="id = %s",
                where_params=(str(session_id),),
            )

            return await self.get_upload_session(session_id)

    async def set_status(
        self,
        session_id: UUID | str,
        status: UploadSessionStatus,
        failure_reason: Optional[str] = None,
        s3_object_key: Optional[str] = None,
        s3_multipart_upload_id: Optional[str] = None,
        completed_at: Optional[datetime] = None,
    ) -> Optional[UploadSession]:
        with span(
            logger,
            "set_upload_session_status",
            {
                "session_id": str(session_id),
                "status": status.value,
            },
        ):
            set_columns = ["status", "updated_at = CURRENT_TIMESTAMP"]
            set_values: List[object] = [status.value]

            if failure_reason is not None:
                set_columns.append("failure_reason")
                set_values.append(failure_reason)

            if s3_object_key is not None:
                set_columns.append("s3_object_key")
                set_values.append(s3_object_key)

            if s3_multipart_upload_id is not None:
                set_columns.append("s3_multipart_upload_id")
                set_values.append(s3_multipart_upload_id)

            resolved_completed_at = completed_at
            if (
                status == UploadSessionStatus.COMPLETED
                and resolved_completed_at is None
            ):
                resolved_completed_at = datetime.now(timezone.utc)

            if resolved_completed_at is not None:
                set_columns.append("completed_at")
                set_values.append(resolved_completed_at)

            await self.db.update(
                table="upload_sessions",
                set_columns=set_columns,
                set_values=set_values,
                where_clause="id = %s",
                where_params=(str(session_id),),
            )

            return await self.get_upload_session(session_id)

    async def expire_stale(self, now: Optional[datetime] = None) -> int:
        now_value = now or datetime.now(timezone.utc)
        with span(
            logger,
            "expire_stale_upload_sessions",
            {"now": now_value.isoformat()},
        ):
            affected = await self.db.execute_commit(
                """
                UPDATE upload_sessions
                SET status = 'expired',
                    updated_at = CURRENT_TIMESTAMP
                WHERE expires_at < %s
                  AND status IN ('initiated', 'uploading', 'processing')
                """,
                (now_value,),
            )
            return int(affected or 0)

    def _row_to_upload_session(self, row) -> UploadSession:
        return UploadSession(
            id=UUID(str(row[0])),
            user_id=row[1],
            filename=row[2],
            mime_type=row[3],
            total_size_bytes=row[4],
            uploaded_size_bytes=row[5],
            status=UploadSessionStatus(str(row[6])),
            s3_multipart_upload_id=row[7],
            s3_object_key=row[8],
            failure_reason=row[9],
            expires_at=row[10],
            completed_at=row[11],
            created_at=row[12],
            updated_at=row[13],
        )
