"""
Kevros Governance Tools for LangChain

Drop-in tools for any LangChain agent to use the Kevros A2A Governance Gateway.

Usage:
    from kevros_tools import get_governance_tools

    tools = get_governance_tools(api_key="kvrs_...")
    agent = initialize_agent(llm, tools, ...)

Or use individual tools:
    from kevros_tools import GovernanceVerifyTool

    tool = GovernanceVerifyTool(api_key="kvrs_...")
    result = tool.invoke({"action_type": "trade", "action_payload": {...}, "agent_id": "bot-1"})
"""

from __future__ import annotations

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

import httpx
from pydantic import BaseModel, Field

GATEWAY_URL = "https://governance.taskhawktech.com"
TIMEOUT = 30.0


def _post(api_key: str, path: str, body: dict) -> str:
    headers = {"X-API-Key": api_key, "Content-Type": "application/json"}
    with httpx.Client(base_url=GATEWAY_URL, headers=headers, timeout=TIMEOUT) as c:
        resp = c.post(path, json=body)
        return resp.text


def _get(api_key: str, path: str) -> str:
    headers = {"X-API-Key": api_key, "Content-Type": "application/json"}
    with httpx.Client(base_url=GATEWAY_URL, headers=headers, timeout=TIMEOUT) as c:
        resp = c.get(path)
        return resp.text


# ---------------------------------------------------------------------------
# Input schemas
# ---------------------------------------------------------------------------

class VerifyInput(BaseModel):
    action_type: str = Field(description="Type of action (e.g., 'trade', 'api_call')")
    action_payload: Dict[str, Any] = Field(description="Action details to verify")
    agent_id: str = Field(description="Your agent identifier")
    policy_context: Optional[Dict[str, Any]] = Field(
        default=None, description="Policy constraints (max_values, forbidden_keys)"
    )


class AttestInput(BaseModel):
    agent_id: str = Field(description="Your agent identifier")
    action_description: str = Field(description="What you did (human-readable)")
    action_payload: Dict[str, Any] = Field(description="Full action details")
    context: Optional[Dict[str, Any]] = Field(default=None, description="Additional context")


class BindInput(BaseModel):
    agent_id: str = Field(description="Your agent identifier")
    intent_type: str = Field(description="Intent category: NAVIGATION, COMMUNICATION, AI_GENERATED, etc.")
    intent_description: str = Field(description="What you intend to accomplish")
    command_payload: Dict[str, Any] = Field(description="Command bound to this intent")
    goal_state: Optional[Dict[str, Any]] = Field(default=None, description="Target state for outcome verification")


class VerifyOutcomeInput(BaseModel):
    agent_id: str = Field(description="Your agent identifier")
    intent_id: str = Field(description="intent_id from bind")
    binding_id: str = Field(description="binding_id from bind")
    actual_state: Dict[str, Any] = Field(description="State after action execution")
    tolerance: float = Field(default=0.1, description="Acceptable deviation (0=exact, 1=any)")


class BundleInput(BaseModel):
    agent_id: str = Field(description="Your agent identifier")
    time_range_start: Optional[str] = Field(default=None, description="ISO 8601 start")
    time_range_end: Optional[str] = Field(default=None, description="ISO 8601 end")


# ---------------------------------------------------------------------------
# LangChain Tool classes (compatible with langchain-core BaseTool)
# ---------------------------------------------------------------------------

try:
    from langchain_core.tools import BaseTool
except ImportError:
    # Fallback: define a minimal base so the module can still be imported
    class BaseTool:  # type: ignore[no-redef]
        name: str = ""
        description: str = ""
        def __init__(self, **kwargs: Any):
            for k, v in kwargs.items():
                setattr(self, k, v)


class GovernanceVerifyTool(BaseTool):
    name: str = "governance_verify"
    description: str = (
        "Verify an action against policy bounds before executing. "
        "Returns ALLOW, CLAMP, or DENY with hash-chained provenance. $0.01/call."
    )
    args_schema: Type[BaseModel] = VerifyInput
    api_key: str = ""

    def _run(self, action_type: str, action_payload: Dict, agent_id: str,
             policy_context: Optional[Dict] = None, **kwargs: Any) -> str:
        body: Dict[str, Any] = {
            "action_type": action_type,
            "action_payload": action_payload,
            "agent_id": agent_id,
        }
        if policy_context:
            body["policy_context"] = policy_context
        return _post(self.api_key, "/governance/verify", body)


class GovernanceAttestTool(BaseTool):
    name: str = "governance_attest"
    description: str = (
        "Create a hash-chained provenance record for an action taken. "
        "Independently verifiable by any third party. $0.02/call."
    )
    args_schema: Type[BaseModel] = AttestInput
    api_key: str = ""

    def _run(self, agent_id: str, action_description: str,
             action_payload: Dict, context: Optional[Dict] = None, **kwargs: Any) -> str:
        body: Dict[str, Any] = {
            "agent_id": agent_id,
            "action_description": action_description,
            "action_payload": action_payload,
        }
        if context:
            body["context"] = context
        return _post(self.api_key, "/governance/attest", body)


class GovernanceBindTool(BaseTool):
    name: str = "governance_bind"
    description: str = (
        "Declare an intent and cryptographically bind it to a command. "
        "Proves intent-command linkage. Use verify_outcome after execution. $0.02/call."
    )
    args_schema: Type[BaseModel] = BindInput
    api_key: str = ""

    def _run(self, agent_id: str, intent_type: str, intent_description: str,
             command_payload: Dict, goal_state: Optional[Dict] = None, **kwargs: Any) -> str:
        body: Dict[str, Any] = {
            "agent_id": agent_id,
            "intent_type": intent_type,
            "intent_description": intent_description,
            "command_payload": command_payload,
            "intent_source": "AI_PLANNER",
        }
        if goal_state:
            body["goal_state"] = goal_state
        return _post(self.api_key, "/governance/bind", body)


class GovernanceVerifyOutcomeTool(BaseTool):
    name: str = "governance_verify_outcome"
    description: str = (
        "Verify that an action achieved its declared intent. "
        "Closes the intent->command->outcome loop. Free (included with bind)."
    )
    args_schema: Type[BaseModel] = VerifyOutcomeInput
    api_key: str = ""

    def _run(self, agent_id: str, intent_id: str, binding_id: str,
             actual_state: Dict, tolerance: float = 0.1, **kwargs: Any) -> str:
        body = {
            "agent_id": agent_id,
            "intent_id": intent_id,
            "binding_id": binding_id,
            "actual_state": actual_state,
            "tolerance": tolerance,
        }
        return _post(self.api_key, "/governance/verify-outcome", body)


class GovernanceBundleTool(BaseTool):
    name: str = "governance_bundle"
    description: str = (
        "Generate a certifier-grade compliance evidence bundle. "
        "Hash-chained provenance, intent proofs, PQC attestations. $0.25/call."
    )
    args_schema: Type[BaseModel] = BundleInput
    api_key: str = ""

    def _run(self, agent_id: str, time_range_start: Optional[str] = None,
             time_range_end: Optional[str] = None, **kwargs: Any) -> str:
        body: Dict[str, Any] = {
            "agent_id": agent_id,
            "include_intent_chains": True,
            "include_pqc_signatures": True,
            "include_verification_instructions": True,
        }
        if time_range_start:
            body["time_range_start"] = time_range_start
        if time_range_end:
            body["time_range_end"] = time_range_end
        return _post(self.api_key, "/governance/bundle", body)


# ---------------------------------------------------------------------------
# Convenience: get all tools at once
# ---------------------------------------------------------------------------

def get_governance_tools(api_key: str) -> List[BaseTool]:
    """
    Get all Kevros governance tools for a LangChain agent.

    Usage:
        tools = get_governance_tools(api_key="kvrs_...")
        agent = initialize_agent(llm, tools, ...)
    """
    return [
        GovernanceVerifyTool(api_key=api_key),
        GovernanceAttestTool(api_key=api_key),
        GovernanceBindTool(api_key=api_key),
        GovernanceVerifyOutcomeTool(api_key=api_key),
        GovernanceBundleTool(api_key=api_key),
    ]
