"""XAI Orchestrator - Main entry point for the XAI Suite.

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

This module provides the XAIOrchestrator class which manages the XAI process,
from loading explainers to generating explanations with optional IA integration.
"""

from typing import Any, Optional, Union
from pathlib import Path

from xai_suite.explainers.shap_explainer import ShapExplainer
from xai_suite.explainers.lime_explainer import LimeExplainer
from xai_suite.ia.client import IAClient, IAClientError
from xai_suite.utils.markdown_exporter import export_to_markdown


class XAIOrchestrator:
    """
    Manages the XAI process, from loading explainers to generating explanations.

    The orchestrator provides a unified interface for:
    - Adding and managing multiple explanation methods (SHAP, LIME)
    
    - Generating local explanations for model predictions
    - Optionally integrating with the IA service for narrative generation
      and actionable insights
    """

    AVAILABLE_EXPLAINERS = {
        "shap": ShapExplainer,
        "lime": LimeExplainer,
    }

    def __init__(self, model, training_data):
        """
        Initialize the XAI Orchestrator.

        Args:
            model: A trained machine learning model.
            training_data (pd.DataFrame): The data used to train the model.
        """
        if not hasattr(training_data, "columns"):
            raise ValueError(
                "training_data must be a pandas DataFrame with column names."
            )

        self.model = model
        self.training_data = training_data
        self.explainers = {}
        self._ia_client: Optional[IAClient] = None
        self._ia_context: Optional[dict[str, Any]] = None
        self._model_type: Optional[str] = None
        self._domain: str = "generic"
        self._use_v2_api: bool = True

    def enable_ia(
        self,
        endpoint: str,
        api_key: str,
        context: Optional[dict[str, Any]] = None,
        model_type: Optional[str] = None,
        domain: str = "generic",
        timeout: float = 30.0,
        verify_ssl: bool = True,
        use_v2_api: bool = True,
    ):
        """
        Enable Intelligence Amplification for narrative generation.

        Args:
            endpoint: Base URL of the IA service (e.g., 'https://xai.cybernemo.eu')
            api_key: API key for authentication
            context: Optional domain context for narrative generation
            model_type: Optional model type hint (e.g., 'classification', 'anomaly_detection')
            domain: Domain context identifier (e.g., 'generic', 'cybersecurity')
            timeout: Request timeout in seconds
            verify_ssl: Whether to verify SSL certificates
            use_v2_api: Whether to use the v2 API (includes insights and compliance)
        """
        self._ia_client = IAClient(
            endpoint=endpoint,
            api_key=api_key,
            timeout=timeout,
            verify_ssl=verify_ssl,
            domain=domain,
        )
        self._ia_context = context
        self._model_type = model_type
        self._domain = domain
        self._use_v2_api = use_v2_api
        print(f"IA enabled. Connected to {endpoint} (domain: {domain})")

    def disable_ia(self):
        """Disable Intelligence Amplification and close the client."""
        if self._ia_client is not None:
            self._ia_client.close()
            self._ia_client = None
            self._ia_context = None
            self._model_type = None
            self._domain = "generic"
            print("IA disabled.")

    @property
    def ia_enabled(self) -> bool:
        """Check if IA is enabled."""
        return self._ia_client is not None

    def set_domain(self, domain: str) -> None:
        """
        Set the domain context for IA explanations.

        Args:
            domain: Domain identifier (e.g., 'generic', 'cybersecurity')
        """
        self._domain = domain
        if self._ia_client is not None:
            self._ia_client.domain = domain

    def add_explainer(self, name: str, config: Optional[dict] = None):
        """
        Initializes and adds an explainer to the orchestrator.

        Args:
            name: The name of the explainer (e.g., 'shap', 'lime').
            config: Configuration for the explainer. Defaults to None.
        """
        name = name.lower()
        if name not in self.AVAILABLE_EXPLAINERS:
            raise ValueError(
                f"Explainer '{name}' not supported. "
                f"Available: {list(self.AVAILABLE_EXPLAINERS.keys())}"
            )

        if config is None:
            config = {}

        explainer_class = self.AVAILABLE_EXPLAINERS[name]
        self.explainers[name] = explainer_class(
            self.model, self.training_data, **config
        )
        print(f"Successfully added '{name}' explainer.")

    def remove_explainer(self, name: str):
        """
        Remove an explainer from the orchestrator.

        Args:
            name: The name of the explainer to remove.
        """
        name = name.lower()
        if name in self.explainers:
            del self.explainers[name]
            print(f"Removed '{name}' explainer.")
        else:
            print(f"Explainer '{name}' not found.")

    def explain(
        self,
        instance,
        methods: Optional[list[str]] = None,
        include_narrative: bool = True,
        include_insights: bool = True,
        include_compliance: bool = True,
        narrative_context: Optional[dict[str, Any]] = None,
    ) -> dict[str, Any]:
        """
        Generates local explanations for a given instance using specified methods.

        Args:
            instance (pd.DataFrame): A single-row DataFrame for the instance to explain.
            methods: A list of explainer names to use. If None, uses all added explainers.
            include_narrative: Whether to generate narrative (requires IA to be enabled).
            include_insights: Whether to generate insights (v2 API only).
            include_compliance: Whether to include compliance info (v2 API only).
            narrative_context: Optional additional context for this specific explanation.

        Returns:
            dict: A dictionary where keys are explainer names and values are their
                explanations. If IA is enabled, may include:
                - 'narrative': Generated narrative text and metadata
                - 'insights': List of actionable insights (v2 API)
                - 'compliance': EU AI Act compliance information (v2 API)
        """
        if methods is None:
            methods = list(self.explainers.keys())

        results = {}
        for method in methods:
            method = method.lower()
            if method not in self.explainers:
                print(f"Warning: Explainer '{method}' not initialized. Skipping.")
                continue

            print(f"\n--- Generating explanation with {method.upper()} ---")
            results[method] = self.explainers[method].explain_local(instance)

        # Generate narrative and insights if IA is enabled
        if self._ia_client is not None and include_narrative and results:
            ia_results = self._generate_ia_response(
                results,
                narrative_context,
                include_insights,
                include_compliance,
            )
            results.update(ia_results)

        return results

    def _generate_ia_response(
        self,
        explanations: dict[str, dict[str, Any]],
        additional_context: Optional[dict[str, Any]] = None,
        include_insights: bool = True,
        include_compliance: bool = True,
    ) -> dict[str, Any]:
        """
        Generate IA response (narrative, insights, compliance) from explanations.

        Args:
            explanations: Dictionary of explanations keyed by method name.
            additional_context: Optional additional context for this request.
            include_insights: Whether to request insights.
            include_compliance: Whether to request compliance info.

        Returns:
            Dictionary containing narrative and optionally insights/compliance.
        """
        # Merge contexts
        context = {}
        if self._ia_context:
            context.update(self._ia_context)
        if additional_context:
            context.update(additional_context)

        # Use the first available explanation (prefer SHAP)
        primary_method = "shap" if "shap" in explanations else list(explanations.keys())[0]
        primary_explanation = explanations[primary_method]

        print("\n--- Generating IA response ---")

        try:
            if self._use_v2_api:
                # Use v2 API for narrative + insights + compliance
                response = self._ia_client.explain(
                    explanation=primary_explanation,
                    model_type=self._model_type,
                    domain=self._domain,
                    domain_context=context if context else None,
                    include_insights=include_insights,
                    include_compliance=include_compliance,
                )
                return {
                    "narrative": {
                        "text": response.narrative.text,
                        "method": response.narrative.method,
                        "domain": response.narrative.domain,
                        "timestamp": response.narrative.timestamp.isoformat(),
                    },
                    "insights": [
                        {
                            "type": i.type,
                            "priority": i.priority,
                            "title": i.title,
                            "description": i.description,
                            "action": i.action,
                            "related_features": i.related_features,
                        }
                        for i in response.insights
                    ],
                    "compliance": {
                        "audit_id": response.compliance.audit_id,
                        "timestamp": response.compliance.timestamp.isoformat(),
                        "applicable_articles": [
                            {
                                "article": a.article,
                                "title": a.title,
                                "relevance": a.relevance,
                            }
                            for a in response.compliance.applicable_articles
                        ],
                    },
                }
            else:
                # Use v1 API (legacy)
                response = self._ia_client.generate_narrative(
                    explanation=primary_explanation,
                    context=context if context else None,
                    model_type=self._model_type,
                )
                return {
                    "narrative": {
                        "text": response.narrative,
                        "compliance_id": response.compliance_id,
                        "timestamp": response.timestamp.isoformat(),
                        "method": response.method,
                    }
                }

        except IAClientError as e:
            print(f"Warning: Failed to generate IA response: {e.message}")
            return {
                "narrative": {
                    "text": None,
                    "error": e.message,
                    "compliance_id": None,
                }
            }

    def export_markdown(
        self,
        result: dict[str, Any],
        output_path: Union[str, Path],
        **kwargs,
    ) -> str:
        """
        Export explanation result to a Markdown file.

        This is a convenience method that wraps export_to_markdown().

        Args:
            result: The explanation result dictionary from explain().
            output_path: Path for the output Markdown file.
            **kwargs: Additional arguments passed to export_to_markdown().
                See export_to_markdown() for available options including:
                - title: Report title
                - top_n_features: Limit to top N features
                - image_format: 'svg' or 'png'
                - include_shap, include_lime, include_narrative, etc.

        Returns:
            Path to the generated Markdown file.

        Example:
            >>> result = orchestrator.explain(instance)
            >>> orchestrator.export_markdown(result, "report.md", top_n_features=10)
            'report.md'
        """
        return export_to_markdown(result, str(output_path), **kwargs)

    def __enter__(self):
        return self

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