"""
Asynchronous RPC API implementation for Blue Server.
"""

import asyncio
import hashlib
import logging
import os
import socket
import ssl
import struct
from asyncio import StreamReader, StreamWriter
from tempfile import gettempdir
from uuid import uuid4

from ecmind_blue_client.rpc import (
    DEFAULT_CA_DATA,
    JobParameter,
    JobRequestFile,
    JobResponseFile,
    JobResponseFileType,
    JobResult,
    Jobs,
    JobValueTypes,
    RpcConnection,
)
from ecmind_blue_client.rpc._blue_protocol_exception import BlueProtocolException
from ecmind_blue_client.rpc.api import (
    append_inferred_params,
    create_request_bin,
    deserialize_file_header,
    deserialize_job_errors,
    deserialize_job_header,
    deserialize_job_parameters,
    serialize_file_footer,
    serialize_file_header,
    serialize_job_header,
)

logger = logging.getLogger(__name__)


class RpcAsyncConnection(RpcConnection):
    """
    Asynchronous implementation of AbstractConnection representing a TCP/IP connection to the Blue Server.
    """

    def __init__(self, stream_reader: StreamReader, stream_writer: StreamWriter, hostname: str, port: int) -> None:
        super().__init__(hostname, port)

        self.reader = stream_reader
        self.writer = stream_writer


async def async_connect(
    host: str, port: int = 4000, use_ssl: bool = True, connect_timeout: int = 10, cadata: str = DEFAULT_CA_DATA
) -> RpcAsyncConnection:
    """
    Create an async connection to the enaio server and return an :class:`RpcAsyncConnection`.

    Args:
        host (str): The hostname or IP address of the enaio server.
        port (int): The port number of the enaio server (default is 4000).
        use_ssl (bool): Whether to use SSL/TLS encryption for the connection (default is True).
        cadata (str): The CA certificate chain as PEM-formatted string for SSL/TLS authentication
            (default is the built-in server certificate).
        connect_timeout (int): The timeout for the connection attempt in seconds (default is 10).

    Returns:
        RpcAsyncConnection: The established async connection to the enaio server.

    Raises:
        TimeoutError: If the connection attempt times out.
        ConnectionError: If the connection fails (DNS resolution failure, connection refused, etc.).
        ssl.SSLError: If there is an SSL/TLS handshake error during the connection attempt.
    """
    try:
        if use_ssl:
            logger.debug("Using SSL/TLS for connection to %s:%d", host, port)
            ssl_context = ssl.create_default_context(purpose=ssl.Purpose.SERVER_AUTH, cadata=cadata)
            ssl_context.check_hostname = False

            future = asyncio.open_connection(host, port, ssl=ssl_context)
            reader, writer = await asyncio.wait_for(future, timeout=connect_timeout)
            return RpcAsyncConnection(reader, writer, host, port)

        logger.debug("Using plain TCP/IP for connection to %s:%d", host, port)
        future = asyncio.open_connection(host, port)
        reader, writer = await asyncio.wait_for(future, timeout=connect_timeout)
        return RpcAsyncConnection(reader, writer, host, port)

    except asyncio.TimeoutError as exc:
        logger.error("Connection to %s:%d timed out after %d seconds", host, port, connect_timeout)
        raise TimeoutError(f"Connection to {host}:{port} timed out after {connect_timeout} seconds") from exc
    except socket.gaierror as exc:
        logger.error("Failed to resolve host %s: %s", host, exc)
        raise ConnectionError(f"Failed to connect to host {host}:{port}: {exc}") from exc
    except ConnectionRefusedError as exc:
        logger.error("Connection to %s:%d was refused: %s", host, port, exc)
        raise ConnectionError(f"Connection to {host}:{port} was refused: {exc}") from exc
    except ssl.SSLError:
        logger.error("SSL/TLS error while connecting to %s:%d", host, port)
        raise
    except Exception as exc:
        logger.error("Unexpected error while connecting to %s:%d: %s", host, port, exc)
        raise ConnectionError(f"Unexpected error while connecting to {host}:{port}: {exc}") from exc


async def async_call_job(
    connection: RpcAsyncConnection,
    name: str | Jobs,
    parameters: list[JobParameter],
    request_files: list[JobRequestFile] | None = None,
    file_cache_byte_limit: int = 3355443,
    **params: JobValueTypes,
) -> JobResult:
    """
    Execute a job on the enaio server using the provided connection, parameters and files.

    Args:
        connection (RpcAsyncConnection): The async connection to the enaio server.
        name (str | Jobs): The name of the job to execute, either as a string or a :class:`Jobs` enum value.
        parameters (list[JobParameter]): A list of :class:`JobParameter` objects for the job's input parameters.
        request_files (list[JobRequestFile] | None): Files to upload for the job (optional).
        file_cache_byte_limit (int): Response files below this byte threshold are kept in memory as a
            ``bytearray``; larger files are written to a temporary path on disk (default is 3355443 bytes).
        **params (JobValueTypes): Additional keyword parameters whose :class:`JobParameterTypes` is inferred
            automatically from the Python type (``bool`` → ``BOOLEAN``, ``int`` → ``INTEGER``,
            ``float`` → ``DOUBLE``, ``datetime`` → ``DATE_TIME``, ``str`` → ``STRING``,
            ``bytes`` → ``BASE64``). Duplicate names with ``parameters`` raise :exc:`ValueError`.

    Returns:
        JobResult: The result returned by the enaio server.

    Raises:
        ValueError: If a keyword parameter name duplicates an entry in ``parameters``, or if its type
            cannot be inferred.
        BlueProtocolException: If the server response does not conform to the expected binary protocol.
    """

    # Infer the params from their types and append them to the parameters list
    append_inferred_params(parameters, params)

    # ----------------------------------------------------------------------
    # Start write request
    # ----------------------------------------------------------------------

    logger.debug("Starting to call the job '%s' on the server %s:%d", name, connection.hostname, connection.port)

    # create request bin based on the provided parameters and files
    request_bin = create_request_bin(name, parameters, request_files if request_files else [])

    # calculate sha1 digest of the request bin
    write_digest = hashlib.sha1()
    write_digest.update(request_bin)

    # create header for the request
    headers_bin = serialize_job_header(len(request_bin))

    logger.debug("Sending job request to the server %s:%d", connection.hostname, connection.port)
    connection.writer.write(headers_bin)
    connection.writer.write(request_bin)

    # Send the files if any
    if request_files:
        for request_file in request_files:
            logger.debug("Sending file to the server %s:%d", connection.hostname, connection.port)
            file_header_bin = serialize_file_header(request_file.size, request_file.extension)
            connection.writer.write(file_header_bin)
            write_digest.update(file_header_bin)

            with request_file.open() as file_stream:
                data = file_stream.read(4096)
                while data:
                    connection.writer.write(data)
                    write_digest.update(data)
                    data = file_stream.read(4096)

            footer_bin = serialize_file_footer()
            connection.writer.write(footer_bin)
            write_digest.update(footer_bin)

    # Send the digest of all data sent
    connection.writer.write(write_digest.digest())
    await connection.writer.drain()
    logger.debug("Sent all request data, waiting for server's response...")

    # ----------------------------------------------------------------------
    # Start read response
    # ----------------------------------------------------------------------

    # read response header
    logger.debug("Reading server's response header...")
    (_protocol, parameter_length, _version, _compression) = deserialize_job_header(await connection.reader.read(20))
    parameter_length = parameter_length - 20

    # Read the rest of the response
    logger.debug("Reading server's response data...")
    response_body = b""
    while len(response_body) < parameter_length:
        response_body += await connection.reader.read(parameter_length - len(response_body))

    # start to create a digest of the received data
    response_digest = hashlib.sha1()
    response_digest.update(response_body)

    logger.debug("Parse server response data...")

    if response_body[0:1].decode("ascii") != "R":
        raise BlueProtocolException("R-byte must be present in the response body")

    response_body = response_body[1:]

    # Read the internal parameters
    [internal_data_length] = struct.unpack(">I", response_body[:4])
    response_body = response_body[4:]

    internal_parameters = deserialize_job_parameters(response_body[:internal_data_length])
    response_body = response_body[internal_data_length:]

    # Read the parameters
    [data_length] = struct.unpack(">I", response_body[:4])
    response_body = response_body[4:]

    parameter = deserialize_job_parameters(response_body[:data_length])
    response_body = response_body[data_length:]

    # Read the errors
    [error_length] = struct.unpack(">I", response_body[:4])
    response_body = response_body[4:]

    errors = deserialize_job_errors(response_body[:error_length])
    response_body = response_body[error_length:]

    # Read the files
    response_files: list[JobResponseFile] = []

    streams_parameter = next((obj for obj in internal_parameters if obj.name == "streams"), None)

    if streams_parameter and isinstance(streams_parameter.value, int):
        logger.debug("Receive file from server")
        streams = streams_parameter.value
        for _ in range(streams):
            header_raw = await connection.reader.read(32)
            response_digest.update(header_raw)
            [file_length, extension] = deserialize_file_header(header_raw)
            to_file = file_length >= file_cache_byte_limit
            remainder = file_length
            file_name = f"ecmind_{str(uuid4())}.{extension}"
            buffer_size = 4096

            if to_file:
                file_path = os.path.join(gettempdir(), file_name)
                with open(file_path, "wb") as file_pointer:
                    while remainder > 0:
                        file_part: bytes = await connection.reader.read(min(remainder, buffer_size))
                        response_digest.update(file_part)
                        remainder -= len(file_part)
                        file_pointer.write(file_part)
                result_file = JobResponseFile(
                    result_file_type=JobResponseFileType.FILE_PATH,
                    file_name=file_name,
                    file_path=file_path,
                )
            else:
                byte_array = bytearray()
                while remainder > 0:
                    file_part = await connection.reader.read(min(remainder, buffer_size))
                    response_digest.update(file_part)
                    remainder -= len(file_part)
                    byte_array += bytearray(file_part)
                result_file = JobResponseFile(
                    result_file_type=JobResponseFileType.BYTE_ARRAY,
                    file_name=file_name,
                    byte_array=bytes(byte_array),
                )

            response_files.append(result_file)

            footer_raw = await connection.reader.read(20)
            response_digest.update(footer_raw)

    # validate the response digest
    response_digest_received = response_digest.digest()
    response_digest_expected = await connection.reader.read(20)
    if response_digest_received != response_digest_expected:
        raise BlueProtocolException(
            f"Received response digest {response_digest_received} does not match expected response digest {response_digest_expected}"
        )

    logger.debug("Successfully received job response")
    return JobResult(internal_parameters, parameter, response_files, errors)
