"""Transport utilities for asynchronous JSON-RPC client.

Author: Vasiliy Zdanovskiy
email: vasilyvz@gmail.com
"""

from __future__ import annotations

import os
from pathlib import Path
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast

if TYPE_CHECKING:
    from mcp_proxy_adapter.client.jsonrpc_client.schema_generator import (
        SchemaRequestGenerator,
    )

import httpx

from mcp_proxy_adapter.client.jsonrpc_client.exceptions import (
    TransferChecksumMismatchClientError,
    TransferClientError,
    TransferCompressionClientError,
    TransferExpiredClientError,
    TransferRangeNotSatisfiableClientError,
    TransferResumeConflictClientError,
    TransferSessionNotFoundClientError,
    TransferTooLargeClientError,
)
from mcp_proxy_adapter.core.ssl_utils import SSLUtils

PATH_UPLOADS = "/api/transfer/uploads"
PATH_UPLOAD_BY_ID = "/api/transfer/uploads/{transfer_id}"
PATH_UPLOAD_CHUNKS = "/api/transfer/uploads/{transfer_id}/chunks"
PATH_UPLOAD_COMPLETE = "/api/transfer/uploads/{transfer_id}/complete"
PATH_DOWNLOADS = "/api/transfer/downloads"
PATH_DOWNLOAD_BY_ID = "/api/transfer/downloads/{transfer_id}"
PATH_DOWNLOAD_CHUNKS = "/api/transfer/downloads/{transfer_id}/chunks"
HDR_TRANSFER_ID = "X-Transfer-Id"
HDR_TRANSFER_OFFSET = "X-Transfer-Offset"
HDR_TRANSFER_COMPLETED = "X-Transfer-Completed"
HDR_TRANSFER_STATUS = "X-Transfer-Status"
HDR_TRANSFER_WIRE_SIZE_BYTES = "X-Transfer-Wire-Size-Bytes"
HDR_TRANSFER_PLAINTEXT_SIZE_BYTES = "X-Transfer-Plaintext-Size-Bytes"
CONTENT_TYPE_OCTET = "application/octet-stream"
JSON_SUCCESS_FALSE_KEY = "success"
JSON_ERROR_TYPE_KEY = "error_type"
JSON_MESSAGE_KEY = "message"
JSON_DETAILS_KEY = "details"


class JsonRpcTransport:
    """Base transport class providing HTTP primitives."""

    def __init__(
        self,
        protocol: str = "http",
        host: str = "127.0.0.1",
        port: int = 8080,
        token_header: Optional[str] = None,
        token: Optional[str] = None,
        cert: Optional[str] = None,
        key: Optional[str] = None,
        ca: Optional[str] = None,
        check_hostname: bool = False,
        timeout: Optional[float] = None,
    ) -> None:
        """
        Initialize JSON-RPC transport.

        Args:
            protocol: Transport protocol (http, https, mtls)
            host: Server hostname
            port: Server port
            token_header: Header name for authentication token
            token: Authentication token value
            cert: Path to client certificate file
            key: Path to client private key file
            ca: Path to CA certificate file
            check_hostname: Whether to verify hostname in SSL connections
            timeout: HTTP client timeout in seconds. If None, uses value from
                MCP_PROXY_ADAPTER_HTTP_TIMEOUT environment variable, or defaults to 30.0
        """
        scheme = "https" if protocol in ("https", "mtls") else "http"
        self.base_url = f"{scheme}://{host}:{port}"

        self.headers: Dict[str, str] = {"Content-Type": "application/json"}
        if token_header and token:
            self.headers[token_header] = token

        self.verify: Union[bool, str] = True
        self.cert: Optional[Tuple[str, str]] = None
        self._check_hostname = check_hostname

        if protocol in ("https", "mtls"):
            if cert and key:
                self.cert = (str(Path(cert)), str(Path(key)))
            if ca:
                self.verify = str(Path(ca))
            else:
                self.verify = False

        # Determine timeout: parameter > environment variable > default
        if timeout is not None:
            self.timeout = float(timeout)
        else:
            env_timeout = os.getenv("MCP_PROXY_ADAPTER_HTTP_TIMEOUT")
            if env_timeout:
                try:
                    self.timeout = float(env_timeout)
                except (ValueError, TypeError):
                    self.timeout = 30.0
            else:
                self.timeout = 30.0
        self._client: Optional[httpx.AsyncClient] = None
        self._protocol = protocol
        self._host = host
        self._port = port
        self._ca = str(Path(ca)) if ca else None

    @property
    def ws_url(self) -> str:
        """WebSocket base URL for /ws endpoint (e.g. ws://host:port/ws or wss://host:port/ws)."""
        scheme = "wss" if self.base_url.startswith("https") else "ws"
        base = self.base_url.replace("https://", "").replace("http://", "")
        return f"{scheme}://{base.rstrip('/')}/ws"

    def get_ws_ssl_context(self) -> Optional[Any]:
        """
        Return SSL context for WebSocket (wss) or None for ws.

        Uses same client cert/key and CA as HTTP transport.
        """
        if self.base_url.startswith("http://"):
            return None
        from mcp_proxy_adapter.core.ssl_utils import SSLUtils

        ca_cert = None
        if isinstance(self.verify, str):
            ca_cert = self.verify
        elif self._ca:
            ca_cert = self._ca
        client_cert = self.cert[0] if self.cert and len(self.cert) >= 1 else None
        client_key = self.cert[1] if self.cert and len(self.cert) >= 2 else None
        verify = bool(self.verify) if not isinstance(self.verify, str) else True
        ctx = SSLUtils.create_client_ssl_context(
            ca_cert=ca_cert,
            client_cert=client_cert,
            client_key=client_key,
            verify=verify,
            check_hostname=self._check_hostname,
        )
        # Ensure hostname check is applied (framework may ignore config)
        ctx.check_hostname = self._check_hostname
        return ctx

    async def _get_client(self) -> httpx.AsyncClient:
        """Return cached async HTTP client or create it lazily."""

        if self._client is None:
            # If verify is a path to CA cert, create SSL context with hostname check disabled
            verify = self.verify
            if isinstance(verify, str):
                verify_path = Path(verify)
                if verify_path.exists():
                    # Extract client cert and key from self.cert tuple if available
                    client_cert = None
                    client_key = None
                    if (
                        self.cert
                        and isinstance(self.cert, tuple)
                        and len(self.cert) == 2
                    ):
                        client_cert = self.cert[0]
                        client_key = self.cert[1]

                    ssl_context = SSLUtils.create_client_ssl_context(
                        ca_cert=str(verify_path),
                        client_cert=client_cert,
                        client_key=client_key,
                        verify=True,
                        check_hostname=self._check_hostname,
                    )
                    # Ensure check_hostname is set correctly (SSLManager may not apply it)
                    if hasattr(ssl_context, "check_hostname"):
                        ssl_context.check_hostname = self._check_hostname
                    verify = ssl_context
                else:
                    # Path doesn't exist, but it's a string - treat as False to avoid errors
                    verify = False

            self._client = httpx.AsyncClient(
                verify=verify,
                cert=self.cert,
                timeout=self.timeout,
            )
        return self._client

    async def close(self) -> None:
        """Close underlying HTTPX client."""

        if self._client:
            await self._client.aclose()
            self._client = None

    async def health(self) -> Dict[str, Any]:
        """Fetch health information from service."""

        client = await self._get_client()
        response = await client.get(f"{self.base_url}/health", headers=self.headers)
        response.raise_for_status()
        return cast(Dict[str, Any], response.json())

    async def get_openapi_schema(self) -> Dict[str, Any]:
        """Fetch OpenAPI schema from service."""

        client = await self._get_client()
        response = await client.get(
            f"{self.base_url}/openapi.json", headers=self.headers
        )
        response.raise_for_status()
        return cast(Dict[str, Any], response.json())

    async def jsonrpc_batch_call(
        self,
        requests: List[Dict[str, Any]],
    ) -> List[Dict[str, Any]]:
        """
        Execute batch JSON-RPC requests.

        Args:
            requests: List of JSON-RPC request dictionaries, each containing:
                - method: Method name
                - params: Method parameters (optional)
                - id: Request ID (optional)

        Returns:
            List of JSON-RPC responses in the same order as requests

        Example:
            requests = [
                {"jsonrpc": "2.0", "method": "echo", "params": {"message": "Hello"}, "id": 1},
                {"jsonrpc": "2.0", "method": "help", "params": {}, "id": 2},
            ]
            responses = await client.jsonrpc_batch_call(requests)
        """
        client = await self._get_client()
        response = await client.post(
            f"{self.base_url}/api/jsonrpc/batch",
            json=requests,
            headers=self.headers,
        )
        response.raise_for_status()
        return cast(List[Dict[str, Any]], response.json())

    async def cmd_call(
        self,
        command: str,
        params: Optional[Dict[str, Any]] = None,
        validate: bool = False,
    ) -> Dict[str, Any]:
        """
        Execute command via /cmd endpoint.

        This method uses the /cmd endpoint which supports the CommandRequest schema
        with oneOf discriminator for different commands. The schema includes detailed
        parameter validation for each command.

        Args:
            command: Command name to execute
            params: Optional command parameters (specific to command)
            validate: If True, validate request against OpenAPI schema before sending

        Returns:
            Command execution result

        Raises:
            RuntimeError: If command execution fails or validation fails
        """
        # Validate against OpenAPI schema if requested
        if validate:
            schema = await self.get_openapi_schema()
            from mcp_proxy_adapter.client.jsonrpc_client.openapi_validator import (
                OpenAPIValidator,
            )

            validator = OpenAPIValidator(schema)
            is_valid, error_msg = validator.validate_command_request(command, params)
            if not is_valid:
                raise RuntimeError(f"Validation error: {error_msg}")

        payload: Dict[str, Any] = {
            "command": command,
        }
        if params:
            payload["params"] = params

        client = await self._get_client()
        response = await client.post(
            f"{self.base_url}/cmd",
            json=payload,
            headers=self.headers,
        )
        response.raise_for_status()
        result = cast(Dict[str, Any], response.json())

        # Check for error in response
        if "error" in result:
            error = result["error"]
            message = error.get("message", "Unknown error")
            code = error.get("code", -1)
            raise RuntimeError(f"Command error: {message} (code: {code})")

        return cast(Dict[str, Any], result.get("result", {}))

    def get_schema_generator(self) -> SchemaRequestGenerator:
        """
        Get schema request generator instance.

        This method requires the schema to be loaded first. Use get_openapi_schema()
        to load the schema, then create a generator.

        Returns:
            SchemaRequestGenerator instance

        Note:
            This is a synchronous method, but it requires the schema to be loaded
            asynchronously first. Consider using get_schema_generator_async() instead.
        """
        # This will fail if schema is not loaded - user should use async version
        raise RuntimeError(
            "Schema must be loaded first. Use get_schema_generator_async() instead, "
            "or load schema with await get_openapi_schema() first."
        )

    async def get_schema_generator_async(self) -> SchemaRequestGenerator:
        """
        Get schema request generator instance with automatic schema loading.

        Returns:
            SchemaRequestGenerator instance
        """
        from mcp_proxy_adapter.client.jsonrpc_client.schema_generator import (
            SchemaRequestGenerator,
        )

        schema = await self.get_openapi_schema()
        return SchemaRequestGenerator(schema)

    async def jsonrpc_call(
        self,
        method: str,
        params: Dict[str, Any],
        req_id: int = 1,
    ) -> Dict[str, Any]:
        """Perform JSON-RPC request and return raw response payload."""

        payload: Dict[str, Any] = {
            "jsonrpc": "2.0",
            "method": method,
            "params": params,
            "id": req_id,
        }
        client = await self._get_client()
        response = await client.post(
            f"{self.base_url}/api/jsonrpc",
            json=payload,
            headers=self.headers,
        )
        response.raise_for_status()
        return cast(Dict[str, Any], response.json())

    def _transfer_absolute_url(self, path: str) -> str:
        return f"{self.base_url.rstrip('/')}{path}"

    def _transfer_headers_octet_stream(self) -> Dict[str, str]:
        headers = dict(self.headers)
        headers["Content-Type"] = CONTENT_TYPE_OCTET
        return headers

    def _raise_for_transfer_http_error(self, response: httpx.Response) -> None:
        if response.is_success:
            return
        try:
            body_any = response.json()
        except Exception:
            response.raise_for_status()
        else:
            if (
                isinstance(body_any, dict)
                and body_any.get(JSON_SUCCESS_FALSE_KEY) is False
            ):
                JsonRpcTransport.raise_for_transfer_error(
                    response.status_code, cast(dict[str, object], body_any)
                )
        response.raise_for_status()

    @staticmethod
    def raise_for_transfer_error(status_code: int, payload: dict[str, object]) -> None:
        etype = payload.get(JSON_ERROR_TYPE_KEY)
        etype_str = etype if isinstance(etype, str) else str(etype)
        msg = str(payload.get(JSON_MESSAGE_KEY) or "")
        raw_details = payload.get(JSON_DETAILS_KEY, {})
        if not isinstance(raw_details, dict):
            raw_details = {}
        details = cast(dict[str, object], raw_details)

        if etype_str == "TransferSessionNotFoundError":
            raise TransferSessionNotFoundClientError(msg, details)
        if etype_str == "TransferResumeConflictError":
            raise TransferResumeConflictClientError(msg, details)
        if etype_str == "TransferChecksumMismatchError":
            raise TransferChecksumMismatchClientError(msg, details)
        if etype_str == "TransferCompressionError":
            raise TransferCompressionClientError(msg, details)
        if etype_str == "TransferTooLargeError":
            raise TransferTooLargeClientError(msg, details)
        if etype_str == "TransferRangeNotSatisfiableError":
            raise TransferRangeNotSatisfiableClientError(msg, details)
        if etype_str == "TransferExpiredError":
            raise TransferExpiredClientError(msg, details)

        detail_msg = f"{etype_str}: {msg}" if msg else etype_str
        raise TransferClientError(f"HTTP {status_code}: {detail_msg}", details)

    async def create_upload_session(
        self,
        filename: str,
        size_bytes: int,
        checksum_algorithm: str,
        checksum_value: str,
        compression: str,
        chunk_size: int | None = None,
    ) -> dict[str, object]:
        """POST /api/transfer/uploads — JSON body; return parsed JSON object."""
        body: dict[str, object] = {
            "filename": filename,
            "size_bytes": size_bytes,
            "checksum_algorithm": checksum_algorithm,
            "checksum_value": checksum_value,
            "compression": compression,
        }
        if chunk_size is not None:
            body["chunk_size"] = chunk_size
        client = await self._get_client()
        response = await client.post(
            self._transfer_absolute_url(PATH_UPLOADS),
            json=body,
            headers=self.headers,
        )
        self._raise_for_transfer_http_error(response)
        return cast(dict[str, object], response.json())

    async def get_upload_session(self, transfer_id: str) -> dict[str, object]:
        """GET /api/transfer/uploads/{transfer_id} — return parsed JSON object."""
        path = PATH_UPLOAD_BY_ID.format(transfer_id=transfer_id)
        client = await self._get_client()
        response = await client.get(
            self._transfer_absolute_url(path),
            headers=self.headers,
        )
        self._raise_for_transfer_http_error(response)
        return cast(dict[str, object], response.json())

    async def upload_chunk(
        self, transfer_id: str, offset: int, chunk_bytes: bytes
    ) -> dict[str, object]:
        """PUT /api/transfer/uploads/{transfer_id}/chunks — raw body; return receipt JSON."""
        path = PATH_UPLOAD_CHUNKS.format(transfer_id=transfer_id)
        url = f"{self._transfer_absolute_url(path)}?offset={offset}"
        client = await self._get_client()
        response = await client.put(
            url,
            content=chunk_bytes,
            headers=self._transfer_headers_octet_stream(),
        )
        self._raise_for_transfer_http_error(response)
        return cast(dict[str, object], response.json())

    async def complete_upload_session(self, transfer_id: str) -> dict[str, object]:
        """POST /api/transfer/uploads/{transfer_id}/complete — return parsed JSON receipt."""
        path = PATH_UPLOAD_COMPLETE.format(transfer_id=transfer_id)
        client = await self._get_client()
        response = await client.post(
            self._transfer_absolute_url(path),
            headers=self.headers,
        )
        self._raise_for_transfer_http_error(response)
        return cast(dict[str, object], response.json())

    async def create_download_session(
        self,
        source_path: str,
        filename: str,
        checksum_algorithm: str = "sha256",
        compression: str = "identity",
        chunk_size: int | None = None,
    ) -> dict[str, object]:
        """POST /api/transfer/downloads — JSON body; return parsed JSON object."""
        body: dict[str, object] = {
            "source_path": source_path,
            "filename": filename,
            "checksum_algorithm": checksum_algorithm,
            "compression": compression,
        }
        if chunk_size is not None:
            body["chunk_size"] = chunk_size
        client = await self._get_client()
        response = await client.post(
            self._transfer_absolute_url(PATH_DOWNLOADS),
            json=body,
            headers=self.headers,
        )
        self._raise_for_transfer_http_error(response)
        return cast(dict[str, object], response.json())

    async def get_download_session(self, transfer_id: str) -> dict[str, object]:
        """GET /api/transfer/downloads/{transfer_id} — return parsed JSON object."""
        path = PATH_DOWNLOAD_BY_ID.format(transfer_id=transfer_id)
        client = await self._get_client()
        response = await client.get(
            self._transfer_absolute_url(path),
            headers=self.headers,
        )
        self._raise_for_transfer_http_error(response)
        return cast(dict[str, object], response.json())

    async def download_chunk(
        self, transfer_id: str, offset: int, limit: int
    ) -> tuple[bytes, dict[str, object]]:
        """GET chunk bytes and header-derived metadata for a download session."""
        path = PATH_DOWNLOAD_CHUNKS.format(transfer_id=transfer_id)
        client = await self._get_client()
        response = await client.get(
            self._transfer_absolute_url(path),
            params={"offset": offset, "limit": limit},
            headers=self.headers,
        )
        if not response.is_success:
            try:
                body_any = response.json()
            except Exception:
                response.raise_for_status()
            else:
                if (
                    isinstance(body_any, dict)
                    and body_any.get(JSON_SUCCESS_FALSE_KEY) is False
                ):
                    JsonRpcTransport.raise_for_transfer_error(
                        response.status_code, cast(dict[str, object], body_any)
                    )
            response.raise_for_status()

        raw_id = response.headers.get(HDR_TRANSFER_ID)
        if not raw_id:
            raise TransferClientError("Missing X-Transfer-Id header")
        raw_off = response.headers.get(HDR_TRANSFER_OFFSET)
        if raw_off is None or str(raw_off).strip() == "":
            raise TransferClientError("Missing X-Transfer-Offset header")
        try:
            off = int(str(raw_off).strip())
        except ValueError as exc:
            raise TransferClientError(
                f"Invalid X-Transfer-Offset header: {raw_off!r}"
            ) from exc

        raw_completed = response.headers.get(HDR_TRANSFER_COMPLETED)
        if raw_completed is None:
            raise TransferClientError("Missing X-Transfer-Completed header")
        completed_lc = str(raw_completed).strip().lower()
        if completed_lc == "true":
            completed = True
        elif completed_lc == "false":
            completed = False
        else:
            raise TransferClientError(
                f"Invalid X-Transfer-Completed header: {raw_completed!r}"
            )

        raw_status = response.headers.get(HDR_TRANSFER_STATUS)
        if raw_status is None:
            raise TransferClientError("Missing X-Transfer-Status header")
        status_val = str(raw_status)
        if status_val == "":
            raise TransferClientError("Missing X-Transfer-Status header")

        metadata: dict[str, object] = {
            "transfer_id": str(raw_id),
            "offset": off,
            "completed": completed,
            "status": status_val,
        }
        raw_wire = response.headers.get(HDR_TRANSFER_WIRE_SIZE_BYTES)
        if raw_wire is not None and str(raw_wire).strip() != "":
            try:
                metadata["wire_size_bytes"] = int(str(raw_wire).strip())
            except ValueError as exc:
                raise TransferClientError(
                    f"Invalid {HDR_TRANSFER_WIRE_SIZE_BYTES} header: {raw_wire!r}"
                ) from exc
        raw_plain = response.headers.get(HDR_TRANSFER_PLAINTEXT_SIZE_BYTES)
        if raw_plain is not None and str(raw_plain).strip() != "":
            try:
                metadata["plaintext_size_bytes"] = int(str(raw_plain).strip())
            except ValueError as exc:
                raise TransferClientError(
                    f"Invalid {HDR_TRANSFER_PLAINTEXT_SIZE_BYTES} header: {raw_plain!r}"
                ) from exc

        return (response.content, metadata)

    async def download_all_chunks_wire(
        self, transfer_id: str, limit: int
    ) -> tuple[bytes, dict[str, object]]:
        """Fetch the full download byte stream by advancing offset from headers.

        Chunk bodies are raw octets (for ``compression=gzip`` on the session, the
        concatenation is one gzip stream; decompress once after the loop).
        """
        parts: list[bytes] = []
        offset = 0
        last_meta: dict[str, object] = {}
        while True:
            chunk, last_meta = await self.download_chunk(transfer_id, offset, limit)
            parts.append(chunk)
            if last_meta.get("completed") is True:
                break
            raw_next = last_meta["offset"]
            if not isinstance(raw_next, int):
                raise TransferClientError(f"Invalid offset metadata: {raw_next!r}")
            offset = raw_next
        return (b"".join(parts), last_meta)

    def _extract_result(self, response: Dict[str, Any]) -> Dict[str, Any]:
        """Extract ``result`` part from JSON-RPC reply raising on error."""

        if "error" in response:
            error = response["error"]
            message = error.get("message", "Unknown error")
            code = error.get("code", -1)
            raise RuntimeError(f"JSON-RPC error: {message} (code: {code})")
        result_data = response.get("result", {})
        return cast(Dict[str, Any], result_data)
