from google.adk.agents.llm_agent import Agent
from kybernis import KybernisSDK

from .prompt import SYSTEM_PROMPT
from .sub_agents import validator_tool, policy_interpreter_tool, executor_tool, audit_logger_tool
from .sub_agents.executor import tools as executor_runtime_tools
import hashlib
import json
import os


def _truthy(name: str, default: bool = False) -> bool:
    raw = os.getenv(name, "").strip().lower()
    if not raw:
        return default
    return raw in {"1", "true", "yes", "on"}


def _json_env(name: str) -> dict:
    """
    Reads a JSON object from `.env` (or process env).
    User-provided input:
    - What: JSON config values from your app.
    - Where: `sdk/python/refund_demo/.env` in this demo, or deployment env vars.
    - Effect: becomes context sent to Kybernis drift guard.
    """
    raw = os.getenv(name, "").strip()
    if not raw:
        return {}
    try:
        loaded = json.loads(raw)
    except Exception:
        return {}
    return loaded if isinstance(loaded, dict) else {}


def _sha256_json(value) -> str:
    blob = json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=False).encode("utf-8")
    return "sha256:" + hashlib.sha256(blob).hexdigest()


def _extract_chat_signature(tool_context) -> str:
    """
    Builds a deterministic hash from ADK session state.
    User-provided input:
    - What: conversation/session state managed by ADK.
    - Where: `tool_context.invocation_context.session.state`.
    - Effect: lets Kybernis detect chat-context drift without storing full chat.
    """
    inv = getattr(tool_context, "invocation_context", None)
    if inv is None:
        return ""
    session = getattr(inv, "session", None)
    if session is None:
        return ""
    state = getattr(session, "state", None)
    if isinstance(state, dict):
        # Keep this lightweight and deterministic.
        subset = {k: state[k] for k in sorted(state.keys())[:20]}
        return _sha256_json(subset)
    return ""


def refund_intent_context_provider(tool_name: str, args: dict, tool_context) -> dict:
    """
    Context provider = app-to-Kybernis adapter for business context.

    Why user defines this:
    - SDK cannot infer your domain semantics.
    - Only your app knows which fields change decision meaning.

    What it does:
    - Collects tool args + domain context.
    - Sends them to server guard (`/infra/intent/guard`) for server-enforced drift checks.
    """
    # User input: pricing flags snapshot.
    # Where: runtime dependency state in executor tools (fallback to .env JSON).
    # Effect: semantic changes (for example true->false) can trigger drift guard.
    pricing_flags = {}
    payment_id = str((args or {}).get("payment_id", "")).strip()
    if payment_id:
        try:
            pricing_flags = executor_runtime_tools.current_pricing_flags(payment_id=payment_id)
        except Exception:
            pricing_flags = {}
    if not pricing_flags:
        pricing_flags = _json_env("REFUND_PRICING_FLAGS_JSON")
    # User input: order metadata snapshot.
    # Where: order API response; demo uses REFUND_ORDER_METADATA_JSON in .env.
    # Effect: included in semantic context used by server-side drift hashing.
    order_metadata = _json_env("REFUND_ORDER_METADATA_JSON")
    # User input: business policy version label (you define this).
    # Where: your policy repo/store release tag, put in REFUND_POLICY_VERSION.
    # Effect: guard ties decisions to policy label for replay/forensics.
    # Note: this is different from workflow_version_id (server-managed workflow graph version).
    policy_version = os.getenv("REFUND_POLICY_VERSION", "refund_policy:v1")

    # Derived dependency signature from user context.
    # Effect: captures upstream dependency meaning changes.
    dependency_material = {
        "pricing_flags": pricing_flags,
        "order_metadata_keys": sorted(order_metadata.keys()),
    }
    dependency_signature = _sha256_json(dependency_material)
    # Derived chat signature from ADK runtime context.
    chat_sig = _extract_chat_signature(tool_context)

    # User input payload sent to server; server computes field semantic hashes.
    semantic_fields = {
        "tool": tool_name,
        "order_metadata": order_metadata,
        "pricing_flags": pricing_flags,
        "chat_context_signature": chat_sig,
        "tool_args": args,
    }
    # Observable execution state (no vendor internals).
    model_label = os.getenv("KYB_MODEL_LABEL", "gemini-2.5-flash").strip()
    provider_label = os.getenv("KYB_PROVIDER_LABEL", "google").strip()
    prompt_template_version = os.getenv("KYB_PROMPT_TEMPLATE_VERSION", "refund_root:v1").strip()
    tool_schema_version = os.getenv("KYB_TOOL_SCHEMA_VERSION", "refund_tools:v1").strip()
    try:
        decoding_params = json.loads(os.getenv("KYB_DECODING_PARAMS_JSON", "{}"))
    except Exception:
        decoding_params = {}
    if not isinstance(decoding_params, dict):
        decoding_params = {}
    decoding_params_signature = _sha256_json({"decoding_params": decoding_params})
    model_execution_signature = _sha256_json(
        {
            "model_label": model_label,
            "provider_label": provider_label,
            "prompt_template_version": prompt_template_version,
            "tool_schema_version": tool_schema_version,
            "decoding_params_signature": decoding_params_signature,
        }
    )
    behavior_probe_mode = os.getenv("KYB_BEHAVIOR_PROBE_MODE", "off").strip().lower()
    if behavior_probe_mode not in {"off", "observe", "enforce"}:
        behavior_probe_mode = "off"
    behavior_probe_suite_version = os.getenv("KYB_BEHAVIOR_PROBE_SUITE_VERSION", "").strip()
    behavior_probe_signature = os.getenv("KYB_BEHAVIOR_PROBE_SIGNATURE", "").strip()
    behavior_probe_checked_at = os.getenv("KYB_BEHAVIOR_PROBE_CHECKED_AT", "").strip()
    # User-defined section ordering.
    # Effect: catches context hierarchy/order shifts.
    context_order = [
        "policy_version",
        "order_metadata",
        "pricing_flags",
        "chat_context_signature",
        "tool_args",
    ]
    context_ordering_signature = "->".join(context_order)
    retrieval_params = {
        "policy_source": "policy_store_v1",
        "pricing_source": "runtime_tools_or_env",
        "order_source": "env_json",
        "assembly_mode": "default",
    }
    retrieval_params_signature = _sha256_json({"retrieval_params": retrieval_params})
    evidence_set_signature = _sha256_json(
        {
            "tool": tool_name,
            "policy_version": policy_version,
            "order_metadata": order_metadata,
            "pricing_flags": pricing_flags,
            "tool_args": args,
        }
    )
    evidence_order_signature = _sha256_json({"context_order": context_order})
    authorization_state_vector_signature = _sha256_json(
        {
            "policy_version": policy_version,
            "dependency_signature": dependency_signature,
            "context_ordering_signature": context_ordering_signature,
            "evidence_set_signature": evidence_set_signature,
            "evidence_order_signature": evidence_order_signature,
            "retrieval_params_signature": retrieval_params_signature,
            "model_label": model_label,
            "provider_label": provider_label,
            "prompt_template_version": prompt_template_version,
            "tool_schema_version": tool_schema_version,
            "decoding_params_signature": decoding_params_signature,
            "model_execution_signature": model_execution_signature,
            "behavior_probe_mode": behavior_probe_mode,
            "behavior_probe_suite_version": behavior_probe_suite_version,
            "behavior_probe_signature": behavior_probe_signature,
        }
    )
    # User-provided stable IDs from tool args.
    # Effect: used for stable intent/idempotency keys across retries/replays.
    identity_fields = {k: v for k, v in args.items() if k.endswith("_id") or k in {"id", "order_id", "customer_id", "payment_id"}}
    function_call_id = str(getattr(tool_context, "function_call_id", "") or "").strip()
    is_mutation_tool = tool_name == "issue_refund"

    # Important for demo stability:
    # - Keep idempotency stable only for the mutating tool (`issue_refund`).
    # - For non-mutating tools (validator/policy/audit/lookups), make each tool call
    #   unique to avoid replay suppression across runs before we reach execution.
    intent_key = ""
    if not is_mutation_tool:
        nonce = function_call_id or _sha256_json({"tool": tool_name, "args": args, "chat": chat_sig})[:16]
        intent_key = f"intent:{tool_name}:{nonce}"
        identity_fields = {"function_call_id": nonce}

    return {
        "policy_version": policy_version,
        "dependency_signature": dependency_signature,
        "evidence_set_signature": evidence_set_signature,
        "evidence_order_signature": evidence_order_signature,
        "retrieval_params_signature": retrieval_params_signature,
        "authorization_state_vector_signature": authorization_state_vector_signature,
        "model_label": model_label,
        "provider_label": provider_label,
        "prompt_template_version": prompt_template_version,
        "tool_schema_version": tool_schema_version,
        "decoding_params": decoding_params,
        "decoding_params_signature": decoding_params_signature,
        "model_execution_signature": model_execution_signature,
        "behavior_probe_mode": behavior_probe_mode,
        "behavior_probe_suite_version": behavior_probe_suite_version,
        "behavior_probe_signature": behavior_probe_signature,
        "behavior_probe_checked_at": behavior_probe_checked_at,
        "retrieval_params": retrieval_params,
        "semantic_fields": semantic_fields,
        "context_order": context_order,
        "context_ordering_signature": context_ordering_signature,
        "identity_fields": identity_fields,
        "intent_key": intent_key,
    }


root_agent = Agent(
    model="gemini-2.5-flash",
    name="refund_root",
    description="Customer-facing refund support agent.",
    instruction=SYSTEM_PROMPT,
    tools=[validator_tool, policy_interpreter_tool, executor_tool, audit_logger_tool],
    output_key="refund_root_output",
)
# Default mode is now zero custom provider:
# - set KYB_ENABLE=0 for baseline (no Kybernis enforcement).
# - set KYB_ENABLE=1 for enforced mode.
# - custom provider is ON by default in this demo because it reads live
#   pricing dependency state, which makes Type-B drift visible to guard.
#
# Set KYB_USE_CUSTOM_CONTEXT_PROVIDER=0 to disable the custom provider and use
# SDK built-ins only.
use_custom_provider = _truthy("KYB_USE_CUSTOM_CONTEXT_PROVIDER", default=True)
mandate_preset = os.getenv("KYB_MANDATE_PRESET", "money")
preset_cfg = KybernisSDK.get_mandate_preset(mandate_preset)
default_semantic_paths = list(preset_cfg.get("semantic_paths") or [])
if not default_semantic_paths:
    default_semantic_paths = [
        "tool_name",
        "tool_args.order_id",
        "tool_args.customer_id",
        "tool_args.payment_id",
        "tool_args.amount_cents",
        "tool_args.reason",
        "session_signature",
    ]
default_session_signature_mode = str(preset_cfg.get("session_signature_mode") or "stable_subset")
env_semantic_paths = os.getenv("KYB_SEMANTIC_PATHS", "").strip()
effective_semantic_paths = None if env_semantic_paths else default_semantic_paths
default_context_mode = "provider" if use_custom_provider else "adk_default"

# Kybernis integration point is this file (agent.py) + env vars.
# User-provided env vars:
# - KYB_API_KEY: from Kybernis credentials.
# - KYB_API_BASE: Kybernis server URL.
# - KYB_TENANT_ID: tenant/workspace id from Kybernis.
# - KYB_ENABLE: 1 enables Kybernis wrapping, 0 disables it for baseline run.
# - KYB_DRIFT_GUARD_MODE: "soft_fail" (block) or "warn_only" (log only).
# - KYB_CONTEXT_MODE: "adk_default" (built-in) | "provider" (custom) | "auto".
#   Today, "auto" and "adk_default" are effectively the same for Google ADK.
# - KYB_SESSION_SIGNATURE_MODE: "stable_subset" (default) | "full" | "off".
# - KYB_SEMANTIC_PATHS: optional CSV override for semantic paths.
# - KYB_MANDATE_PRESET: "money" | "infra" | "low_risk" preset helper.
#
# Version terms:
# - REFUND_POLICY_VERSION: your business-policy label (you control this).
# - workflow_version_id: Kybernis workflow graph version (server controls this).
if _truthy("KYB_ENABLE", default=True):
    KybernisSDK(
        root_agent,
        # For drift-enforcement demos, custom provider should usually be enabled.
        intent_context_provider=refund_intent_context_provider if use_custom_provider else None,
        context_mode=os.getenv("KYB_CONTEXT_MODE", default_context_mode),
        mandate_preset=mandate_preset,
        semantic_paths=effective_semantic_paths,
        session_signature_mode=os.getenv("KYB_SESSION_SIGNATURE_MODE", default_session_signature_mode),
        drift_guard_enabled=True,
        drift_guard_mode=os.getenv("KYB_DRIFT_GUARD_MODE", "soft_fail"),
        # Fallback policy label if provider omits one.
        policy_version=os.getenv("REFUND_POLICY_VERSION", "refund_policy:v1"),
    )
