"""HTTP client for the IA (Intelligence Amplification) service.

Copyright 2026 Maggioli SPA
Licensed under Apache 2.0
Author: Costas Vrioni (costas.vrioni@maggioli.gr)

This module provides the IAClient class for communicating with the XAI Suite
IA Service, supporting both v1 and v2 API versions.
"""

from typing import Any, Optional

import httpx

from xai_suite.ia.models import (
    ExplanationRequest,
    ExplanationRequestV2,
    ExplanationResponseV2,
    NarrativeResponse,
    HealthResponse,
    DomainsResponse,
)


class IAClientError(Exception):
    """Exception raised for IA client errors."""

    def __init__(self, message: str, status_code: Optional[int] = None):
        self.message = message
        self.status_code = status_code
        super().__init__(self.message)


class IAClient:
    """
    HTTP client for communicating with the XAI Suite IA Service.

    The IA service provides narrative generation using LLMs and
    EU AI Act compliance logging.

    Supports both v1 API (legacy) and v2 API (enhanced with insights).
    """

    def __init__(
        self,
        endpoint: str,
        api_key: str,
        timeout: float = 30.0,
        verify_ssl: bool = True,
        domain: str = "generic",
    ):
        """
        Initialize the IA client.

        Args:
            endpoint: Base URL of the IA service (e.g., 'https://xai.cybernemo.eu')
            api_key: API key for authentication
            timeout: Request timeout in seconds
            verify_ssl: Whether to verify SSL certificates
            domain: Default domain context for explanations (e.g., 'generic', 'cybersecurity')
        """
        self.endpoint = endpoint.rstrip("/")
        self.api_key = api_key
        self.timeout = timeout
        self.verify_ssl = verify_ssl
        self.domain = domain
        self._client: Optional[httpx.Client] = None

    def _get_client(self) -> httpx.Client:
        """Get or create the HTTP client."""
        if self._client is None:
            self._client = httpx.Client(
                base_url=self.endpoint,
                timeout=self.timeout,
                verify=self.verify_ssl,
                headers={
                    "Authorization": f"Bearer {self.api_key}",
                    "Content-Type": "application/json",
                },
            )
        return self._client

    def close(self):
        """Close the HTTP client."""
        if self._client is not None:
            self._client.close()
            self._client = None

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()

    # =========================================================================
    # Health & Discovery
    # =========================================================================

    def health_check(self) -> HealthResponse:
        """
        Check if the IA service is healthy.

        Returns:
            HealthResponse with service status
        """
        client = self._get_client()
        try:
            response = client.get("/api/v1/health")
            response.raise_for_status()
            return HealthResponse(**response.json())
        except httpx.HTTPStatusError as e:
            raise IAClientError(
                f"Health check failed: {e.response.text}",
                status_code=e.response.status_code,
            )
        except httpx.RequestError as e:
            raise IAClientError(f"Connection error: {str(e)}")

    def list_domains(self) -> DomainsResponse:
        """
        List available domain contexts.

        Returns:
            DomainsResponse with list of available domains
        """
        client = self._get_client()
        try:
            response = client.get("/api/v1/domains")
            response.raise_for_status()
            return DomainsResponse(**response.json())
        except httpx.HTTPStatusError as e:
            raise IAClientError(
                f"Failed to list domains: {e.response.text}",
                status_code=e.response.status_code,
            )
        except httpx.RequestError as e:
            raise IAClientError(f"Connection error: {str(e)}")

    # =========================================================================
    # V1 API (Legacy)
    # =========================================================================

    def generate_narrative(
        self,
        explanation: dict[str, Any],
        context: Optional[dict[str, Any]] = None,
        model_type: Optional[str] = None,
    ) -> NarrativeResponse:
        """
        Send explanation to server and get narrative back (v1 API).

        Args:
            explanation: Dictionary containing explanation data from SHAP/LIME
            context: Optional additional context for narrative generation
            model_type: Optional model type hint (e.g., 'classification', 'anomaly_detection')

        Returns:
            NarrativeResponse with generated narrative and compliance ID
        """
        # Build the request from the explanation dict
        request = ExplanationRequest(
            method=explanation.get("method", "unknown"),
            feature_names=explanation.get("feature_names", []),
            importance_scores=explanation.get("importance_scores", []),
            feature_values=explanation.get("feature_values"),
            base_value=explanation.get("base_value"),
            output_value=explanation.get("output_value"),
            model_type=model_type,
            context=context,
        )

        client = self._get_client()
        try:
            response = client.post(
                "/api/v1/narrative",
                json=request.model_dump(exclude_none=True),
            )
            response.raise_for_status()
            return NarrativeResponse(**response.json())
        except httpx.HTTPStatusError as e:
            raise IAClientError(
                f"Narrative generation failed: {e.response.text}",
                status_code=e.response.status_code,
            )
        except httpx.RequestError as e:
            raise IAClientError(f"Connection error: {str(e)}")

    def generate_narratives_batch(
        self,
        explanations: dict[str, dict[str, Any]],
        context: Optional[dict[str, Any]] = None,
        model_type: Optional[str] = None,
    ) -> dict[str, NarrativeResponse]:
        """
        Generate narratives for multiple explanations (v1 API).

        Args:
            explanations: Dictionary keyed by explainer name with explanation data
            context: Optional additional context for narrative generation
            model_type: Optional model type hint

        Returns:
            Dictionary keyed by explainer name with NarrativeResponse objects
        """
        results = {}
        for name, explanation in explanations.items():
            if name == "narrative":
                continue  # Skip if already contains narrative
            try:
                results[name] = self.generate_narrative(
                    explanation, context=context, model_type=model_type
                )
            except IAClientError as e:
                # Log the error but continue with other explanations
                print(f"Warning: Failed to generate narrative for {name}: {e.message}")
        return results

    # =========================================================================
    # V2 API (Enhanced)
    # =========================================================================

    def explain(
        self,
        explanation: dict[str, Any],
        model_type: Optional[str] = None,
        domain: Optional[str] = None,
        domain_context: Optional[dict[str, Any]] = None,
        include_insights: bool = True,
        include_compliance: bool = True,
    ) -> ExplanationResponseV2:
        """
        Generate explanation with narrative, insights, and compliance info (v2 API).

        This is the recommended method for new integrations. It provides:
        - Human-readable narrative from SHAP/LIME explanations
        - Actionable insights based on domain context
        - EU AI Act compliance documentation

        Args:
            explanation: Dictionary containing explanation data from SHAP/LIME
            model_type: Optional model type hint (classification, anomaly_detection, regression)
            domain: Domain context to use (default: client's default domain)
            domain_context: Additional domain-specific context
            include_insights: Whether to generate actionable insights
            include_compliance: Whether to include EU AI Act compliance references

        Returns:
            ExplanationResponseV2 with narrative, insights, and compliance info
        """
        request = ExplanationRequestV2(
            method=explanation.get("method", "unknown"),
            feature_names=explanation.get("feature_names", []),
            importance_scores=explanation.get("importance_scores", []),
            feature_values=explanation.get("feature_values"),
            base_value=explanation.get("base_value"),
            output_value=explanation.get("output_value"),
            model_type=model_type,
            domain=domain or self.domain,
            domain_context=domain_context,
            include_insights=include_insights,
            include_compliance=include_compliance,
        )

        client = self._get_client()
        try:
            response = client.post(
                "/api/v1/v2/explain",
                json=request.model_dump(exclude_none=True),
            )
            response.raise_for_status()
            return ExplanationResponseV2(**response.json())
        except httpx.HTTPStatusError as e:
            raise IAClientError(
                f"Explanation generation failed: {e.response.text}",
                status_code=e.response.status_code,
            )
        except httpx.RequestError as e:
            raise IAClientError(f"Connection error: {str(e)}")

    def explain_batch(
        self,
        explanations: dict[str, dict[str, Any]],
        model_type: Optional[str] = None,
        domain: Optional[str] = None,
        domain_context: Optional[dict[str, Any]] = None,
        include_insights: bool = True,
        include_compliance: bool = True,
    ) -> dict[str, ExplanationResponseV2]:
        """
        Generate explanations for multiple SHAP/LIME results (v2 API).

        Args:
            explanations: Dictionary keyed by explainer name with explanation data
            model_type: Optional model type hint
            domain: Domain context to use
            domain_context: Additional domain-specific context
            include_insights: Whether to generate actionable insights
            include_compliance: Whether to include EU AI Act compliance references

        Returns:
            Dictionary keyed by explainer name with ExplanationResponseV2 objects
        """
        results = {}
        for name, explanation in explanations.items():
            if name in ("narrative", "insights", "compliance"):
                continue  # Skip reserved keys
            try:
                results[name] = self.explain(
                    explanation=explanation,
                    model_type=model_type,
                    domain=domain,
                    domain_context=domain_context,
                    include_insights=include_insights,
                    include_compliance=include_compliance,
                )
            except IAClientError as e:
                print(f"Warning: Failed to generate explanation for {name}: {e.message}")
        return results
