"""Client for royalty_accounting MySQL database operations.

Manages adjustment batch lifecycle and staging data in the royalty_accounting database.
Handles batch status transitions (pending → validating → complete/error), error tracking,
and loading validated adjustments into the staging table for downstream processing.

The update_batch_status method supports optimistic locking via expected_status parameter,
ensuring atomic state transitions without concurrent modification conflicts.
"""

from __future__ import annotations

import json
from collections.abc import Generator
from contextlib import contextmanager
from pathlib import Path

from pymysql.cursors import Cursor as PyMySQLCursor

from src.connectors.mysql import (
    MySQLConnection,
    handle_mysql_errors,
    load_from_file,
    load_from_s3,
)
from src.constants import AdjustmentStagingSchema
from src.enums import BatchErrorCode
from src.infra.log import logger


class RoyaltyAccountingClient:
    """Client for royalty accounting database operations.

    Provides high-level operations for managing adjustment batches and staging
    data. Handles batch status transitions, error tracking, and staging table
    management with transaction support.
    """

    def __init__(self, conn: MySQLConnection) -> None:
        """Initialize repository.

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

    @handle_mysql_errors
    def delete_batch_from_staging(
        self,
        batch_id: int,
    ) -> int:
        """Clear data for a specific batch.

        This ensures clean slate before loading new data.

        Note: Caller is responsible for committing the transaction.

        Args:
            batch_id: Batch ID to clear data for

        Returns:
            int: Number of rows deleted

        Raises:
            TransientError: If database connection fails
            Exception: For other database errors
        """
        query = """
            DELETE FROM staging_adjustment_detail
            WHERE worksheet_flowthrough_batch_id = %s
        """
        params = [batch_id]

        with self.conn.cursor() as cursor:
            cursor.execute(query, params)
            return cursor.rowcount

    @handle_mysql_errors
    def stage_batch_from_file(self, file_path: str | Path) -> None:
        """Load prepared CSV data into MySQL staging table.

        Clears any existing staging data for the batch, then loads the CSV
        file into the MySQL staging table. This makes the data available for
        downstream processing.

        Uses LOAD DATA LOCAL INFILE from local file path (DEV environment).

        Args:
            file_path: Local file path for DEV environment loading.

        Raises:
            TransientError: If database connection fails.
            Exception: For other errors including file access issues.
        """
        logger.info(f'Staging "{file_path}"')
        table_name = AdjustmentStagingSchema.TABLE
        with self._stage_connection() as cursor:
            row_count = load_from_file(
                cursor,
                file_path,
                table_name,
                columns=AdjustmentStagingSchema.COLUMNS,
            )
        logger.info(f'Loaded {row_count} rows into staging table: {table_name}')

    @handle_mysql_errors
    def stage_batch_from_s3(self, s3_bucket: str, s3_key: str) -> None:
        """Load prepared CSV data into MySQL staging table.

        Clears any existing staging data for the batch, then loads the CSV
        file into the MySQL staging table. This makes the data available for
        downstream processing.

        Uses LOAD DATA FROM S3 (Aurora MySQL).

        Args:
            s3_bucket: S3 bucket containing the prepared CSV file.
            s3_key: S3 key of the prepared CSV file.

        Raises:
            TransientError: If database connection fails.
            Exception: For other errors including S3 access issues.
        """
        logger.info(f'Staging "s3://{s3_bucket}/{s3_key}"')
        table_name = AdjustmentStagingSchema.TABLE
        with self._stage_connection() as cursor:
            row_count = load_from_s3(
                cursor,
                s3_bucket,
                s3_key,
                table_name,
                columns=AdjustmentStagingSchema.COLUMNS,
            )
        logger.info(f'Loaded {row_count} rows into staging table: {table_name}')

    @handle_mysql_errors
    def update_batch_status(
        self,
        batch_id: int,
        status: str,
        expected_status: str | None = None,
        errors: list[BatchErrorCode] | None = None,
    ) -> int:
        """Update batch status.

        Note: Caller is responsible for committing the transaction.

        Args:
            batch_id: Batch ID
            status: New status
            expected_status: Optional status to check against before updating
            errors: Optional list of errors to update

        Returns:
            int: Number of rows affected

        Raises:
            TransientError: If database connection fails
            Exception: For other database errors
        """
        query = """
            UPDATE worksheet_flowthrough_batch
            SET batch_status = %s
        """
        params: list[int | str] = [status]

        if errors:
            query += ', errors = %s'
            params.append(json.dumps([e.value for e in errors]))

        query += ' WHERE worksheet_flowthrough_batch_id = %s'
        params.append(batch_id)

        if expected_status:
            query += ' AND batch_status = %s'
            params.append(expected_status)

        with self.conn.cursor() as cursor:
            return cursor.execute(query, params)

    @contextmanager
    def _stage_connection(self) -> Generator[PyMySQLCursor, None, None]:
        with self.conn.cursor() as cursor:
            try:
                # FK checks disabled for performance. FKs are checked during
                # validation and all involved entities are soft-deleted. Any
                # hard-deletes (for accounts, contracts) are rare manual
                # occurrences, with low impact risk. Alternatively, set to
                # self.env.is_managed to only disable in dev environments.
                cursor.execute('SET SESSION FOREIGN_KEY_CHECKS = 0')
                cursor.execute('SET SESSION UNIQUE_CHECKS = 0')
                yield cursor
            finally:
                cursor.execute('SET SESSION FOREIGN_KEY_CHECKS = 1')
                cursor.execute('SET SESSION UNIQUE_CHECKS = 1')