"""LangChain and LlamaIndex tool integrations for HIPAA Agent.

These wrappers let you drop HIPAA Agent capabilities directly into any
LangChain agent or LlamaIndex workflow.  Both frameworks are **optional
dependencies** — import errors are deferred to instantiation time so the
core SDK never requires them.

LangChain usage::

    from hipaa_agent import HIPAAAgent
    from hipaa_agent.tools import HIPAAAgentToolkit

    client = HIPAAAgent(api_key="...")
    toolkit = HIPAAAgentToolkit(client=client)
    tools = toolkit.get_tools()  # list[BaseTool]

LlamaIndex usage::

    from hipaa_agent import HIPAAAgent
    from hipaa_agent.tools import HIPAAAgentToolSpec

    client = HIPAAAgent(api_key="...")
    spec = HIPAAAgentToolSpec(client=client)
    tools = spec.to_tool_list()
"""

from __future__ import annotations

from typing import Any, Dict, List, Optional, Type

from pydantic import BaseModel, Field

from hipaa_agent.client import HIPAAAgent

# ---------------------------------------------------------------------------
# LangChain integration
# ---------------------------------------------------------------------------

_LANGCHAIN_AVAILABLE = True
try:
    from langchain_core.tools import BaseTool as _LCBaseTool
except ImportError:
    _LANGCHAIN_AVAILABLE = False
    _LCBaseTool = None  # type: ignore[assignment, misc]


def _require_langchain() -> None:
    if not _LANGCHAIN_AVAILABLE:
        raise ImportError(
            "langchain-core is required for LangChain tool integration. "
            "Install it with: pip install 'hipaa-agent[langchain]'"
        )


# --- Pydantic input schemas ------------------------------------------------

class ScanPracticeInput(BaseModel):
    """Input schema for scanning a healthcare practice."""

    npi: str = Field(
        ...,
        description="10-digit National Provider Identifier of the practice to scan.",
    )
    domain: Optional[str] = Field(
        default=None,
        description="Practice website domain. Auto-discovered from NPI if omitted.",
    )


class GetComplianceScoreInput(BaseModel):
    """Input schema for retrieving a practice's compliance score."""

    npi: str = Field(
        ...,
        description="10-digit NPI of the practice.",
    )


class GetBreachInput(BaseModel):
    """Input schema for looking up HHS breach records."""

    entity_name: Optional[str] = Field(
        default=None,
        description="Name of the breached entity to search for.",
    )
    state: Optional[str] = Field(
        default=None,
        description="Two-letter US state code to filter breaches.",
    )
    limit: int = Field(
        default=10,
        description="Maximum number of breach records to return.",
    )


class ValidateWorkflowInput(BaseModel):
    """Input schema for validating a HIPAA compliance workflow."""

    workflow_type: str = Field(
        ...,
        description=(
            "Type of workflow to validate, e.g. 'breach_response', "
            "'access_request', 'sra'."
        ),
    )
    npi: Optional[str] = Field(
        default=None,
        description="Optional NPI for practice-specific validation context.",
    )


class GetBreachProbabilityInput(BaseModel):
    """Input schema for calculating breach probability."""

    npi: str = Field(
        ...,
        description="10-digit NPI of the practice.",
    )
    time_horizon_days: int = Field(
        default=365,
        description="Prediction window in days.",
    )


class CheckVendorInput(BaseModel):
    """Input schema for checking a vendor's HIPAA compliance."""

    vendor_name: str = Field(
        ...,
        description="Name of the vendor to check.",
    )
    vendor_domain: Optional[str] = Field(
        default=None,
        description="Domain of the vendor's website.",
    )


class GetComplianceStateInput(BaseModel):
    """Input schema for retrieving full compliance state."""

    npi: str = Field(
        ...,
        description="10-digit NPI of the practice.",
    )


# --- LangChain Tool classes -------------------------------------------------

if _LANGCHAIN_AVAILABLE:

    class ScanPracticeTool(_LCBaseTool):  # type: ignore[misc]
        """LangChain tool that launches a HIPAA compliance scan."""

        name: str = "hipaa_scan_practice"
        description: str = (
            "Launch a comprehensive HIPAA compliance scan for a healthcare practice. "
            "Requires the practice's 10-digit NPI. Optionally accepts a website domain. "
            "Returns scan ID, grade, score, and finding count."
        )
        args_schema: Type[BaseModel] = ScanPracticeInput
        client: Any = Field(exclude=True)

        class Config:
            arbitrary_types_allowed = True

        def _run(self, npi: str, domain: Optional[str] = None) -> Dict[str, Any]:
            return self.client.scan(npi=npi, domain=domain)

    class GetComplianceScoreTool(_LCBaseTool):  # type: ignore[misc]
        """LangChain tool that retrieves a practice's HIPAA compliance score."""

        name: str = "hipaa_get_compliance_score"
        description: str = (
            "Get the HIPAA Agent Compliance Score for a practice. "
            "Returns the letter grade (A-F), numeric score, and per-category breakdown."
        )
        args_schema: Type[BaseModel] = GetComplianceScoreInput
        client: Any = Field(exclude=True)

        class Config:
            arbitrary_types_allowed = True

        def _run(self, npi: str) -> Dict[str, Any]:
            return self.client.get_compliance_score(npi=npi)

    class GetBreachTool(_LCBaseTool):  # type: ignore[misc]
        """LangChain tool that looks up HHS breach records."""

        name: str = "hipaa_get_breach"
        description: str = (
            "Search the HHS breach database for healthcare data breach records. "
            "Can filter by entity name, US state, or both. Returns breach details "
            "including severity, individuals affected, and breach type."
        )
        args_schema: Type[BaseModel] = GetBreachInput
        client: Any = Field(exclude=True)

        class Config:
            arbitrary_types_allowed = True

        def _run(
            self,
            entity_name: Optional[str] = None,
            state: Optional[str] = None,
            limit: int = 10,
        ) -> Any:
            return self.client.get_breach(
                entity_name=entity_name,
                state=state,
                limit=limit,
            )

    class ValidateWorkflowTool(_LCBaseTool):  # type: ignore[misc]
        """LangChain tool that validates HIPAA compliance workflows."""

        name: str = "hipaa_validate_workflow"
        description: str = (
            "Validate a HIPAA compliance workflow (e.g. breach response, access request, "
            "security risk assessment) for correctness and completeness. "
            "Returns validation status, errors, and warnings."
        )
        args_schema: Type[BaseModel] = ValidateWorkflowInput
        client: Any = Field(exclude=True)

        class Config:
            arbitrary_types_allowed = True

        def _run(
            self,
            workflow_type: str,
            npi: Optional[str] = None,
        ) -> Dict[str, Any]:
            return self.client.validate_workflow(
                workflow_type=workflow_type,
                npi=npi,
            )

    class GetBreachProbabilityTool(_LCBaseTool):  # type: ignore[misc]
        """LangChain tool that calculates breach probability."""

        name: str = "hipaa_get_breach_probability"
        description: str = (
            "Calculate the probability of a data breach for a healthcare practice "
            "based on its current security posture, historical data, and industry trends. "
            "Returns probability score, risk level, and contributing factors."
        )
        args_schema: Type[BaseModel] = GetBreachProbabilityInput
        client: Any = Field(exclude=True)

        class Config:
            arbitrary_types_allowed = True

        def _run(
            self,
            npi: str,
            time_horizon_days: int = 365,
        ) -> Dict[str, Any]:
            return self.client.get_breach_probability(
                npi=npi,
                time_horizon_days=time_horizon_days,
            )

    class CheckVendorTool(_LCBaseTool):  # type: ignore[misc]
        """LangChain tool that checks a vendor's HIPAA compliance posture."""

        name: str = "hipaa_check_vendor"
        description: str = (
            "Check whether a technology vendor is HIPAA-compliant and has a BAA available. "
            "Returns compliance status, BAA availability, and risk level."
        )
        args_schema: Type[BaseModel] = CheckVendorInput
        client: Any = Field(exclude=True)

        class Config:
            arbitrary_types_allowed = True

        def _run(
            self,
            vendor_name: str,
            vendor_domain: Optional[str] = None,
        ) -> Dict[str, Any]:
            return self.client.check_vendor(
                vendor_name=vendor_name,
                vendor_domain=vendor_domain,
            )

    class GetComplianceStateTool(_LCBaseTool):  # type: ignore[misc]
        """LangChain tool that retrieves full compliance state."""

        name: str = "hipaa_get_compliance_state"
        description: str = (
            "Get the complete current HIPAA compliance state for a practice, "
            "including overall status and per-category breakdown. "
            "Useful for understanding the full compliance picture at a point in time."
        )
        args_schema: Type[BaseModel] = GetComplianceStateInput
        client: Any = Field(exclude=True)

        class Config:
            arbitrary_types_allowed = True

        def _run(self, npi: str) -> Dict[str, Any]:
            return self.client.get_compliance_state(npi=npi)

else:
    # Stub classes when langchain is not installed.  They raise on
    # instantiation so static analysis can still see the names.

    class ScanPracticeTool:  # type: ignore[no-redef]
        def __init__(self, *a: Any, **kw: Any) -> None:
            _require_langchain()

    class GetComplianceScoreTool:  # type: ignore[no-redef]
        def __init__(self, *a: Any, **kw: Any) -> None:
            _require_langchain()

    class GetBreachTool:  # type: ignore[no-redef]
        def __init__(self, *a: Any, **kw: Any) -> None:
            _require_langchain()

    class ValidateWorkflowTool:  # type: ignore[no-redef]
        def __init__(self, *a: Any, **kw: Any) -> None:
            _require_langchain()

    class GetBreachProbabilityTool:  # type: ignore[no-redef]
        def __init__(self, *a: Any, **kw: Any) -> None:
            _require_langchain()

    class CheckVendorTool:  # type: ignore[no-redef]
        def __init__(self, *a: Any, **kw: Any) -> None:
            _require_langchain()

    class GetComplianceStateTool:  # type: ignore[no-redef]
        def __init__(self, *a: Any, **kw: Any) -> None:
            _require_langchain()


# --- Toolkit ----------------------------------------------------------------

class HIPAAAgentToolkit:
    """Convenience wrapper that returns all LangChain tools for a client.

    Usage::

        toolkit = HIPAAAgentToolkit(client=client)
        agent = initialize_agent(tools=toolkit.get_tools(), ...)
    """

    def __init__(self, client: HIPAAAgent) -> None:
        _require_langchain()
        self._client = client

    def get_tools(self) -> List[Any]:
        """Return a list of LangChain ``BaseTool`` instances."""
        return [
            ScanPracticeTool(client=self._client),
            GetComplianceScoreTool(client=self._client),
            GetBreachTool(client=self._client),
            ValidateWorkflowTool(client=self._client),
            GetBreachProbabilityTool(client=self._client),
            CheckVendorTool(client=self._client),
            GetComplianceStateTool(client=self._client),
        ]


# ---------------------------------------------------------------------------
# LlamaIndex integration
# ---------------------------------------------------------------------------

_LLAMA_INDEX_AVAILABLE = True
try:
    from llama_index.core.tools import FunctionTool as _LIFunctionTool
    from llama_index.core.tools.tool_spec.base import BaseToolSpec as _LIBaseToolSpec
except ImportError:
    _LLAMA_INDEX_AVAILABLE = False
    _LIFunctionTool = None  # type: ignore[assignment, misc]
    _LIBaseToolSpec = None  # type: ignore[assignment, misc]


def _require_llama_index() -> None:
    if not _LLAMA_INDEX_AVAILABLE:
        raise ImportError(
            "llama-index-core is required for LlamaIndex tool integration. "
            "Install it with: pip install 'hipaa-agent[llama-index]'"
        )


class HIPAAAgentToolSpec:
    """LlamaIndex tool spec that exposes HIPAA Agent capabilities.

    Usage::

        spec = HIPAAAgentToolSpec(client=client)
        tools = spec.to_tool_list()
    """

    spec_functions: List[str] = [
        "scan_practice",
        "get_compliance_score",
        "get_breach",
        "validate_workflow",
        "get_breach_probability",
        "check_vendor",
        "get_compliance_state",
    ]

    def __init__(self, client: HIPAAAgent) -> None:
        _require_llama_index()
        self._client = client

    # --- Tool methods (match spec_functions) --------------------------------

    def scan_practice(
        self,
        npi: str,
        domain: Optional[str] = None,
    ) -> Dict[str, Any]:
        """Launch a comprehensive HIPAA compliance scan for a healthcare practice.

        Args:
            npi: 10-digit National Provider Identifier.
            domain: Practice website domain. Auto-discovered if omitted.
        """
        return self._client.scan(npi=npi, domain=domain)

    def get_compliance_score(self, npi: str) -> Dict[str, Any]:
        """Get the HIPAA Agent Compliance Score for a practice.

        Args:
            npi: 10-digit NPI of the practice.
        """
        return self._client.get_compliance_score(npi=npi)

    def get_breach(
        self,
        entity_name: Optional[str] = None,
        state: Optional[str] = None,
        limit: int = 10,
    ) -> Any:
        """Search the HHS breach database for healthcare data breach records.

        Args:
            entity_name: Name of the breached entity.
            state: Two-letter US state code.
            limit: Max number of records to return.
        """
        return self._client.get_breach(
            entity_name=entity_name,
            state=state,
            limit=limit,
        )

    def validate_workflow(
        self,
        workflow_type: str,
        npi: Optional[str] = None,
    ) -> Dict[str, Any]:
        """Validate a HIPAA compliance workflow for correctness.

        Args:
            workflow_type: E.g. 'breach_response', 'access_request', 'sra'.
            npi: Optional practice NPI for context.
        """
        return self._client.validate_workflow(
            workflow_type=workflow_type,
            npi=npi,
        )

    def get_breach_probability(
        self,
        npi: str,
        time_horizon_days: int = 365,
    ) -> Dict[str, Any]:
        """Calculate breach probability for a practice.

        Args:
            npi: 10-digit NPI.
            time_horizon_days: Prediction window in days.
        """
        return self._client.get_breach_probability(
            npi=npi,
            time_horizon_days=time_horizon_days,
        )

    def check_vendor(
        self,
        vendor_name: str,
        vendor_domain: Optional[str] = None,
    ) -> Dict[str, Any]:
        """Check a vendor's HIPAA compliance posture.

        Args:
            vendor_name: Name of the vendor.
            vendor_domain: Domain of the vendor's website.
        """
        return self._client.check_vendor(
            vendor_name=vendor_name,
            vendor_domain=vendor_domain,
        )

    def get_compliance_state(self, npi: str) -> Dict[str, Any]:
        """Get complete current HIPAA compliance state for a practice.

        Args:
            npi: 10-digit NPI.
        """
        return self._client.get_compliance_state(npi=npi)

    # --- Conversion to LlamaIndex tools ------------------------------------

    def to_tool_list(self) -> List[Any]:
        """Convert all spec functions to a list of LlamaIndex FunctionTools."""
        _require_llama_index()
        tools: List[Any] = []
        for fn_name in self.spec_functions:
            fn = getattr(self, fn_name)
            tools.append(
                _LIFunctionTool.from_defaults(
                    fn=fn,
                    name=f"hipaa_{fn_name}",
                    description=fn.__doc__ or "",
                )
            )
        return tools
