"""Batch session client for managing batch upload sessions."""

from types import TracebackType
from typing import final

from ksef2.core import exceptions
from ksef2.core.protocols import Middleware
from ksef2.domain.models.batch import PartUploadRequest, PreparedBatch
from ksef2.domain.models import BatchSessionState
from ksef2.domain.models.session import (
    SessionInvoicesResponse,
    SessionStatusResponse,
)
from ksef2.endpoints.invoices import InvoicesEndpoints
from ksef2.endpoints.session import SessionEndpoints
from ksef2.infra.mappers.sessions import from_spec as session_from_spec
from ksef2.logging import get_logger

logger = get_logger(__name__)


@final
class BatchSessionClient:
    """Client for managing an active batch session.

    This client is returned when opening a batch session and provides
    methods for managing the session lifecycle.
    """

    def __init__(
        self,
        transport: Middleware,
        state: BatchSessionState,
        *,
        upload_transport: Middleware | None = None,
        prepared_batch: PreparedBatch | None = None,
    ) -> None:
        self._transport = transport
        self._upload_transport = upload_transport or transport
        self._state = state
        self._prepared_batch = prepared_batch
        self._invoice_eps = InvoicesEndpoints(transport)
        self._session_eps = SessionEndpoints(transport)
        self._closed = False

    def _ensure_open(self) -> None:
        if self._closed:
            raise exceptions.KSeFClientClosedError("Session client is closed.")

    @property
    def reference_number(self) -> str:
        """Get the batch session reference number."""
        return self._state.reference_number

    @property
    def access_token(self) -> str:
        """Get the access token for this session."""
        self._ensure_open()
        return self._state.access_token

    @property
    def aes_key(self) -> bytes:
        """Get the AES key for encrypting batch files."""
        self._ensure_open()
        return self._state.get_aes_key_bytes()

    @property
    def iv(self) -> bytes:
        """Get the initialization vector for AES encryption."""
        self._ensure_open()
        return self._state.get_iv_bytes()

    @property
    def part_upload_requests(self) -> list[PartUploadRequest]:
        """Get the upload instructions for each file part."""
        self._ensure_open()
        return self._state.part_upload_requests

    def get_state(self) -> BatchSessionState:
        """Get the serializable state of this batch session.

        The returned state can be serialized to JSON and used later
        to resume the session or access upload URLs.

        Returns:
            BatchSessionState containing all session information.
        """
        return self._state

    def get_status(self) -> SessionStatusResponse:
        """Fetch the current processing state of the batch session."""
        return session_from_spec(
            self._invoice_eps.get_session_status(
                reference_number=self._state.reference_number,
            )
        )

    def list_invoices(
        self,
        *,
        page_size: int = 10,
        continuation_token: str | None = None,
    ) -> SessionInvoicesResponse:
        """Fetch one page of invoices submitted in this batch session."""
        return session_from_spec(
            self._invoice_eps.list_session_invoices(
                reference_number=self._state.reference_number,
                continuation_token=continuation_token,
                pageSize=page_size,
            )
        )

    def list_failed_invoices(
        self,
        *,
        page_size: int = 10,
        continuation_token: str | None = None,
    ) -> SessionInvoicesResponse:
        """Fetch one page of failed invoices from this batch session."""
        return session_from_spec(
            self._invoice_eps.list_failed_session_invoices(
                reference_number=self._state.reference_number,
                continuation_token=continuation_token,
                pageSize=page_size,
            )
        )

    def get_upo(self, *, upo_reference_number: str) -> bytes:
        """Download the collective UPO for the batch session."""
        return self._session_eps.get_session_upo(
            reference_number=self._state.reference_number,
            upo_reference_number=upo_reference_number,
        )

    def upload_parts(self) -> None:
        """Upload the prepared batch parts using the session's presigned URLs.

        Raises:
            KSeFClientClosedError: If the upload window has already been closed.
            KSeFValidationError: If the session has no prepared batch payload attached.
        """
        self._ensure_open()

        if self._prepared_batch is None:
            raise exceptions.KSeFValidationError(
                "Batch session has no prepared batch attached. "
                "Open it through auth.batch_session(prepared_batch=...) "
                "or auth.batch.open_session(prepared_batch=...)."
            )

        upload_requests = {
            request.ordinal_number: request
            for request in self._state.part_upload_requests
        }
        parts = {part.ordinal_number: part for part in self._prepared_batch.parts}

        if set(upload_requests) != set(parts):
            raise exceptions.KSeFValidationError(
                "Prepared parts do not match the batch session upload instructions.",
                upload_ordinals=sorted(upload_requests),
                prepared_ordinals=sorted(parts),
            )

        for ordinal_number in sorted(upload_requests):
            upload_request = upload_requests[ordinal_number]
            part = parts[ordinal_number]
            _ = self._upload_transport.request(
                upload_request.method,
                upload_request.url,
                headers={
                    "Content-Type": "application/octet-stream",
                    **{
                        key: value
                        for key, value in upload_request.headers.items()
                        if value is not None
                    },
                },
                content=part.content,
            )

    def close(self) -> None:
        """Close the batch session and start processing.

        This triggers processing of the invoice batch and generation of UPO
        for valid invoices and a collective UPO for the session.
        """
        if self._closed:
            return

        self._session_eps.close_batch(
            reference_number=self._state.reference_number,
        )
        self._closed = True

    def __enter__(self) -> "BatchSessionClient":
        self._ensure_open()
        return self

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: TracebackType | None,
    ) -> None:
        try:
            self.close()
        except exceptions.KSeFException:
            if exc_type is None:
                raise
            logger.warning(
                "Failed to close batch session",
                reference_number=self._state.reference_number,
                exc_info=True,
            )
