"""Repository for worksheet_flowthrough operations."""

from __future__ import annotations

from pymysql.connections import Connection
from pymysql.cursors import DictCursor

from src.connectors.db import handle_mysql_errors
from src.enums import (
    BatchStatus,
    BatchType,
    StatementPeriodStatus,
)
from src.schemas import (
    AdjustmentBatch,
    FileUpload,
    StatementPeriod,
)


class Repository:
    """Repository for database operations."""

    def __init__(self, conn: Connection[DictCursor]) -> None:
        """Initialize repository.

        Args:
            conn: MySQL database connection
        """
        self.conn = conn

    @handle_mysql_errors
    def create_batch(
        self, statement_period_id: int, file_upload_id: int, user_id: str
    ) -> int:
        """Create batch record in worksheet_flowthrough_batch table.

        Note: Caller is responsible for committing the transaction.

        Args:
            statement_period_id: The statement period ID
            file_upload_id: The source file_upload ID
            user_id: User who created the upload

        Returns:
            int: Created batch_id

        Raises:
            TransientError: If database connection fails
            Exception: For other database errors
        """
        with self.conn.cursor() as cursor:
            cursor.execute(
                """
                INSERT INTO worksheet_flowthrough_batch (
                    statement_period_id,
                    source_file_upload_id,
                    batch_type,
                    batch_status,
                    created_by,
                    last_modified_by
                ) VALUES (%s, %s, %s, %s, %s, %s)
                """,
                (
                    statement_period_id,
                    file_upload_id,
                    BatchType.UPLOAD,
                    BatchStatus.PENDING,
                    user_id,
                    user_id,
                ),
            )

            return cursor.lastrowid

    @handle_mysql_errors
    def get_batch_by_file_upload(self, file_upload_id: int) -> AdjustmentBatch | None:
        """Get the ID for a worksheet_flowthrough_batch record by its source_file_upload_id.

        Args:
            file_upload_id: Primary key of file_upload table

        Returns:
            AdjustmentBatch with batch details or None if not found

        Raises:
            TransientError: If database connection fails (retriable)
            Exception: For other database errors
        """
        with self.conn.cursor() as cursor:
            cursor.execute(
                """
                SELECT
                    worksheet_flowthrough_batch_id AS batch_id,
                    batch_type,
                    batch_status,
                    source_file_upload_id,
                    statement_period_id
                FROM worksheet_flowthrough_batch
                WHERE source_file_upload_id = %s
                    AND deleted_at IS NULL
                LIMIT 1
                """,
                (file_upload_id,),
            )

            result = cursor.fetchone()
            return AdjustmentBatch(**result) if result else None

    @handle_mysql_errors
    def get_current_statement_period(self) -> StatementPeriod | None:
        """Get current statement_period.

        Returns:
            StatementPeriod with statement details or None if not found

        Raises:
            TransientError: If database connection fails (retriable)
            Exception: For other database errors
        """
        with self.conn.cursor() as cursor:
            cursor.execute(
                """
                SELECT
                    statement_period_id,
                    statement_period_name,
                    statement_period_status,
                    statement_month,
                    statement_year
                FROM statement_period
                WHERE statement_period_status = %s
                LIMIT 1
                """,
                (StatementPeriodStatus.CURRENT,),
            )

            result = cursor.fetchone()
            return StatementPeriod(**result) if result else None

    @handle_mysql_errors
    def get_file_upload(self, file_upload_id: int) -> FileUpload | None:
        """Get file_upload record by ID.

        Args:
            file_upload_id: Primary key of file_upload table

        Returns:
            FileUpload with file details or None if not found

        Raises:
            TransientError: If database connection fails (retriable)
            Exception: For other database errors
        """
        with self.conn.cursor() as cursor:
            cursor.execute(
                """
                SELECT
                    fu.file_upload_id,
                    fu.original_file_name,
                    fu.s3_bucket,
                    fu.s3_key,
                    fu.upload_status,
                    fuc.upload_type,
                    fu.created_by
                FROM file_upload AS fu
                    INNER JOIN file_upload_config AS fuc
                        ON fu.file_upload_config_id = fuc.file_upload_config_id
                WHERE fu.file_upload_id = %s
                    AND fu.deleted_at IS NULL
                LIMIT 1
                """,
                (file_upload_id,),
            )

            result = cursor.fetchone()
            return FileUpload(**result) if result else None