"""
Pramana Engine — LangChain Integration

Provides:
  - PramanaVerifier: Callback handler that auto-verifies LLM outputs
  - PramanaGuardChain: Runnable that verifies chain output via pipe operator
  - PramanaToolkit: LangChain Tool for agent-driven verification

All langchain imports are wrapped in try/except so pramana does not
hard-depend on langchain.
"""

import json
import logging
from dataclasses import dataclass, field
from typing import Any, Optional

from pramana.core.engine import Pramana, VerificationResult
from pramana.constraints.agent_safety import verify_llm_output

logger = logging.getLogger("pramana.integrations.langchain")

# ---------------------------------------------------------------------------
# Lazy langchain imports
# ---------------------------------------------------------------------------
try:
    from langchain_core.callbacks import BaseCallbackHandler
except ImportError:
    try:
        from langchain.callbacks.base import BaseCallbackHandler
    except ImportError:
        BaseCallbackHandler = None

try:
    from langchain_core.runnables import RunnableConfig, RunnableLambda
except ImportError:
    try:
        from langchain.schema.runnable import RunnableConfig, RunnableLambda
    except ImportError:
        RunnableConfig = None
        RunnableLambda = None

try:
    from langchain_core.tools import BaseTool
except ImportError:
    try:
        from langchain.tools import BaseTool
    except ImportError:
        BaseTool = None


# ---------------------------------------------------------------------------
# Exception
# ---------------------------------------------------------------------------

class PramanaVerificationError(Exception):
    """Raised when Pramana verification fails on an LLM output."""

    def __init__(self, result: VerificationResult, message: str = ""):
        self.result = result
        msg = message or (
            f"Pramana verification {result.verdict}: "
            f"{'; '.join(result.reasoning)}"
        )
        super().__init__(msg)


# ---------------------------------------------------------------------------
# 1. PramanaVerifier — Callback Handler
# ---------------------------------------------------------------------------

if BaseCallbackHandler is not None:
    class PramanaVerifier(BaseCallbackHandler):
        """
        LangChain callback handler that verifies every LLM output through
        the Pramana engine before it is returned to the caller.

        Usage:
            from pramana.integrations.langchain import PramanaVerifier

            verifier = PramanaVerifier(
                banned_patterns=["password", "ssn"],
                max_length=2000,
                on_fail="raise",   # "raise" | "warn"
            )
            llm = ChatOpenAI(callbacks=[verifier])
        """

        def __init__(
            self,
            constraints: dict = None,
            banned_patterns: list = None,
            required_fields: list = None,
            expected_format: str = None,
            max_length: int = None,
            profile: str = None,
            on_fail: str = "warn",
            **kwargs,
        ):
            super().__init__(**kwargs)
            self.constraints = constraints or {}
            self.banned_patterns = banned_patterns
            self.required_fields = required_fields
            self.expected_format = expected_format
            self.max_length = max_length
            self.profile = profile
            self.on_fail = on_fail  # "raise" or "warn"
            self.last_result: Optional[VerificationResult] = None

        def on_llm_end(self, response, **kwargs):
            """Called when an LLM call finishes. Verifies the output."""
            # Extract text from the response
            text = self._extract_text(response)
            if text is None:
                return

            result = verify_llm_output(
                output=text,
                expected_format=self.expected_format,
                max_length=self.max_length,
                required_fields=self.required_fields,
                banned_patterns=self.banned_patterns,
            )
            self.last_result = result

            if result.verdict in ("FLAGGED", "HALTED"):
                if self.on_fail == "raise":
                    raise PramanaVerificationError(result)
                else:
                    logger.warning(
                        "Pramana verification %s: %s",
                        result.verdict,
                        "; ".join(result.reasoning),
                    )

        @staticmethod
        def _extract_text(response) -> Optional[str]:
            """Pull the generated text out of an LLMResult."""
            try:
                # LLMResult.generations is list[list[Generation]]
                for gen_list in response.generations:
                    for gen in gen_list:
                        if hasattr(gen, "text") and gen.text:
                            return gen.text
                        if hasattr(gen, "message"):
                            content = getattr(gen.message, "content", None)
                            if content:
                                return content
            except (AttributeError, IndexError):
                pass
            return None
else:
    # langchain not installed — provide a stub that errors clearly
    class PramanaVerifier:  # type: ignore[no-redef]
        def __init__(self, *args, **kwargs):
            raise ImportError(
                "LangChain is required for PramanaVerifier. "
                "Install it with: pip install langchain-core"
            )


# ---------------------------------------------------------------------------
# 2. PramanaGuardChain — Runnable for pipe operator
# ---------------------------------------------------------------------------

def _build_guard_fn(
    constraints: dict = None,
    banned_patterns: list = None,
    required_fields: list = None,
    expected_format: str = None,
    max_length: int = None,
    on_fail: str = "raise",
):
    """Return a function suitable for RunnableLambda."""

    def _guard(input_value):
        # Accept string or object with .content
        if isinstance(input_value, str):
            text = input_value
        elif hasattr(input_value, "content"):
            text = input_value.content
        elif isinstance(input_value, dict) and "output" in input_value:
            text = input_value["output"]
        elif isinstance(input_value, dict) and "text" in input_value:
            text = input_value["text"]
        else:
            text = str(input_value)

        result = verify_llm_output(
            output=text,
            expected_format=expected_format,
            max_length=max_length,
            required_fields=required_fields,
            banned_patterns=banned_patterns,
        )

        if result.verdict in ("FLAGGED", "HALTED") and on_fail == "raise":
            raise PramanaVerificationError(result)

        return {
            "output": text,
            "pramana_verdict": result.verdict,
            "pramana_confidence": result.confidence,
            "pramana_reasoning": result.reasoning,
            "pramana_metadata": result.metadata,
        }

    return _guard


def PramanaGuardChain(
    constraints: dict = None,
    banned_patterns: list = None,
    required_fields: list = None,
    expected_format: str = None,
    max_length: int = None,
    on_fail: str = "raise",
):
    """
    Create a Runnable that verifies chain output through Pramana.

    Usage:
        from pramana.integrations.langchain import PramanaGuardChain

        chain = my_llm_chain | PramanaGuardChain(
            banned_patterns=["password"],
            max_length=500,
        )
        result = chain.invoke({"question": "..."})
        # result["pramana_verdict"]  -> "VERIFIED" | "FLAGGED" | "HALTED"
    """
    if RunnableLambda is None:
        raise ImportError(
            "LangChain is required for PramanaGuardChain. "
            "Install it with: pip install langchain-core"
        )

    fn = _build_guard_fn(
        constraints=constraints,
        banned_patterns=banned_patterns,
        required_fields=required_fields,
        expected_format=expected_format,
        max_length=max_length,
        on_fail=on_fail,
    )
    return RunnableLambda(fn)


# Alias for more natural pipe syntax: chain | PramanaGuard(...)
PramanaGuard = PramanaGuardChain


# ---------------------------------------------------------------------------
# 3. PramanaToolkit — LangChain Tool for agents
# ---------------------------------------------------------------------------

if BaseTool is not None:
    class PramanaToolkit(BaseTool):
        """
        LangChain Tool that agents can call to verify data against constraints.

        Tool name: pramana_verify

        Input: JSON string with keys:
            - data (dict): The data to verify
            - constraints (dict, optional): Constraints to check against
            - banned_patterns (list, optional): Patterns to reject
            - required_fields (list, optional): Fields that must be present
            - max_length (int, optional): Max output length
            - expected_format (str, optional): e.g. "json"

        Returns: JSON string with verification result.

        Usage:
            from pramana.integrations.langchain import PramanaToolkit
            tools = [PramanaToolkit()]
            agent = create_react_agent(llm, tools)
        """

        name: str = "pramana_verify"
        description: str = (
            "Verify data through the Pramana epistemological engine. "
            "Input must be a JSON string with 'data' (required) and optional "
            "'constraints', 'banned_patterns', 'required_fields', "
            "'max_length', 'expected_format'. "
            "Returns a JSON object with verdict, confidence, and reasoning."
        )

        def _run(self, tool_input: str, **kwargs) -> str:
            """Synchronous verification."""
            try:
                params = json.loads(tool_input)
            except (json.JSONDecodeError, TypeError):
                return json.dumps({
                    "error": "Invalid JSON input",
                    "verdict": "HALTED",
                })

            data = params.get("data")
            if data is None:
                return json.dumps({
                    "error": "Missing 'data' field in input",
                    "verdict": "HALTED",
                })

            # If data is a string, run LLM output verification
            if isinstance(data, str):
                result = verify_llm_output(
                    output=data,
                    expected_format=params.get("expected_format"),
                    max_length=params.get("max_length"),
                    required_fields=params.get("required_fields"),
                    banned_patterns=params.get("banned_patterns"),
                )
            else:
                # Data is a dict — run constraint verification
                engine = Pramana()
                constraints = params.get("constraints", {})
                result = engine.verify_quick(
                    data=data, constraints=constraints
                )

            return json.dumps(result.to_dict(), default=str)

        async def _arun(self, tool_input: str, **kwargs) -> str:
            """Async — delegates to sync since Pramana is CPU-bound."""
            return self._run(tool_input, **kwargs)
else:
    class PramanaToolkit:  # type: ignore[no-redef]
        def __init__(self, *args, **kwargs):
            raise ImportError(
                "LangChain is required for PramanaToolkit. "
                "Install it with: pip install langchain-core"
            )
