from __future__ import annotations

import hashlib
import importlib.util
import json
import os
import sys
import time
import uuid
from pathlib import Path
from typing import Any, Dict

# Make sdk/python importable when script is executed directly.
_HERE = Path(__file__).resolve()
_SDK_PYTHON = _HERE.parents[1]
if str(_SDK_PYTHON) not in sys.path:
    sys.path.insert(0, str(_SDK_PYTHON))

from kybernis.client import KybernisClient


def _load_executor_tools_module():
    tools_path = _HERE.parent / "sub_agents" / "executor" / "tools.py"
    spec = importlib.util.spec_from_file_location("refund_demo_executor_tools", str(tools_path))
    if spec is None or spec.loader is None:
        raise RuntimeError(f"unable to load executor tools module from {tools_path}")
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


executor_tools = _load_executor_tools_module()


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


def _state_file() -> Path:
    raw = os.getenv("KYB_VIDEO_STATE_FILE", "").strip()
    if raw:
        return Path(raw)
    return Path(__file__).resolve().parent / ".payments_video_state.json"


def _refund_key(customer_id: str, payment_id: str, amount_cents: int, reason: str) -> str:
    return f"{customer_id}|{payment_id}|{amount_cents}|{reason}"


def _seed_state(path: Path, *, customer_id: str, payment_id: str, amount_cents: int, reason: str) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    seed = {
        "fail_once_keys": {
            # mark fail-once as already consumed so first call succeeds
            _refund_key(customer_id, payment_id, amount_cents, reason): True
        },
        "payments": {
            customer_id: [
                {"payment_id": "pay_87421", "amount_cents": amount_cents, "status": "succeeded", "created_day": "yesterday"},
                {"payment_id": payment_id, "amount_cents": amount_cents, "status": "succeeded", "created_day": "yesterday"},
            ]
        },
    }
    path.write_text(json.dumps(seed), encoding="utf-8")


def _build_intent(
    *,
    tool_name: str,
    policy_version: str,
    pricing_flags: Dict[str, Any],
    order_metadata: Dict[str, Any],
    tool_args: Dict[str, Any],
    context_order: list[str] | None = None,
    retrieval_params: Dict[str, Any] | None = None,
) -> Dict[str, Any]:
    if context_order is None:
        context_order = ["policy_version", "order_metadata", "pricing_flags", "tool_args"]
    if retrieval_params is None:
        retrieval_params = {
            "retriever": "default",
            "reranker": "v1",
            "top_k": 4,
        }
    dep_material = {
        "policy_version": policy_version,
        "pricing_flags": pricing_flags,
        "order_metadata_keys": sorted(order_metadata.keys()),
    }
    semantic_fields = {
        "order_metadata": order_metadata,
        "pricing_flags": pricing_flags,
        "tool_args": tool_args,
    }
    context_ordering_signature = "->".join(context_order)
    evidence_set_signature = _sha256_json({"semantic_fields": semantic_fields})
    evidence_order_signature = _sha256_json({"context_order": context_order})
    retrieval_params_signature = _sha256_json({"retrieval_params": retrieval_params})
    authorization_state_vector_signature = _sha256_json(
        {
            "policy_version": policy_version,
            "dependency_signature": _sha256_json(dep_material),
            "context_ordering_signature": context_ordering_signature,
            "evidence_set_signature": evidence_set_signature,
            "evidence_order_signature": evidence_order_signature,
            "retrieval_params_signature": retrieval_params_signature,
        }
    )
    return {
        "tool_name": tool_name,
        "policy_version": policy_version,
        "dependency_signature": _sha256_json(dep_material),
        "context_ordering_signature": context_ordering_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,
        "identity_fields": {
            "customer_id": tool_args.get("customer_id"),
            "payment_id": tool_args.get("payment_id"),
            "order_id": tool_args.get("order_id"),
        },
        "semantic_fields": semantic_fields,
    }


def _run_refund_once(tool_args: Dict[str, Any]) -> Dict[str, Any]:
    return executor_tools.issue_refund(
        customer_id=str(tool_args["customer_id"]),
        payment_id=str(tool_args["payment_id"]),
        amount_cents=int(tool_args["amount_cents"]),
        reason=str(tool_args["reason"]),
    )


def _run_before_without_kybernis(
    *,
    tool_args: Dict[str, Any],
    pricing_flags_a: Dict[str, Any],
    pricing_flags_b: Dict[str, Any],
) -> Dict[str, Any]:
    # Baseline path intentionally ignores pricing flag drift at execution gate.
    # This simulates "no Kybernis guard".
    first = _run_refund_once(tool_args)
    second = _run_refund_once(tool_args)
    return {
        "pricing_flags_attempt_1": pricing_flags_a,
        "pricing_flags_attempt_2": pricing_flags_b,
        "attempt_1_result": first,
        "attempt_2_result": second,
        "attempt_2_executed": True,
    }


def _run_after_with_kybernis(
    *,
    client: KybernisClient,
    run_id: str,
    idempotency_key: str,
    policy_version: str,
    order_metadata: Dict[str, Any],
    pricing_flags_a: Dict[str, Any],
    pricing_flags_b: Dict[str, Any],
    context_order_a: list[str],
    context_order_b: list[str],
    retrieval_params_a: Dict[str, Any],
    retrieval_params_b: Dict[str, Any],
    tool_args: Dict[str, Any],
    requested_mode: str,
) -> Dict[str, Any]:
    intent1 = _build_intent(
        tool_name="issue_refund",
        policy_version=policy_version,
        pricing_flags=pricing_flags_a,
        order_metadata=order_metadata,
        tool_args=tool_args,
        context_order=context_order_a,
        retrieval_params=retrieval_params_a,
    )
    guard1 = client.intent_guard(
        run_id=run_id,
        idempotency_key=idempotency_key,
        intent=intent1,
        original_payload={"intent": intent1, "tool_args": tool_args},
        requested_mode=requested_mode,
    )
    if bool(guard1.get("block_execute")):
        return {
            "error": "attempt_1_blocked_unexpectedly",
            "guard_1": guard1,
        }

    attempt1_result = _run_refund_once(tool_args)
    mutation_id = str(guard1.get("mutation_id") or "")
    runtime_intent_hash = str(guard1.get("runtime_intent_hash") or "")
    if mutation_id:
        if str(attempt1_result.get("status")) == "succeeded":
            client.intent_commit(
                mutation_id=mutation_id,
                runtime_intent_hash=runtime_intent_hash,
                receipt={"tool": "issue_refund", "status": "ok", "provider_result": attempt1_result},
                provider="refund_demo_video",
            )
        else:
            client.intent_fail(mutation_id=mutation_id, reason=str(attempt1_result.get("message") or "failed"))

    intent2 = _build_intent(
        tool_name="issue_refund",
        policy_version=policy_version,
        pricing_flags=pricing_flags_b,
        order_metadata=order_metadata,
        tool_args=tool_args,
        context_order=context_order_b,
        retrieval_params=retrieval_params_b,
    )
    guard2 = client.intent_guard(
        run_id=run_id,
        idempotency_key=idempotency_key,
        intent=intent2,
        original_payload={"intent": intent2, "tool_args": tool_args},
        requested_mode=requested_mode,
    )
    blocked = bool(guard2.get("block_execute"))
    attempt2_result: Dict[str, Any]
    if blocked:
        attempt2_result = {
            "status": "blocked_by_kybernis",
            "reason": str(guard2.get("reason") or ""),
        }
    else:
        attempt2_result = _run_refund_once(tool_args)

    return {
        "guard_1": guard1,
        "attempt_1_result": attempt1_result,
        "guard_2": guard2,
        "attempt_2_result": attempt2_result,
        "attempt_2_executed": not blocked,
    }


def main() -> int:
    customer_id = os.getenv("KYB_VIDEO_CUSTOMER_ID", "cust_4829")
    payment_id = os.getenv("KYB_VIDEO_PAYMENT_ID", "pay_87422")
    amount_cents = int(os.getenv("KYB_VIDEO_AMOUNT_CENTS", "12000"))
    reason = os.getenv("KYB_VIDEO_REASON", "duplicate_charge")
    order_id = os.getenv("KYB_VIDEO_ORDER_ID", "ord_video_1001")
    policy_version = os.getenv("REFUND_POLICY_VERSION", "refund_policy:v1")
    requested_mode = os.getenv("KYB_DRIFT_GUARD_MODE", "soft_fail")
    workflow_version_id = os.getenv("KYB_WORKFLOW_VERSION_ID", "").strip()
    drift_kind = os.getenv("KYB_VIDEO_DRIFT_KIND", "b1").strip().lower()
    scenario_nonce = uuid.uuid4().hex[:8]

    tool_args: Dict[str, Any] = {
        "customer_id": customer_id,
        "payment_id": payment_id,
        "amount_cents": amount_cents,
        "reason": reason,
        "order_id": order_id,
    }
    order_metadata = {
        "order_id": order_id,
        "customer_id": customer_id,
        "channel": "support_chat",
    }
    if drift_kind == "b2":
        # B2 demo: same business args and same dependency snapshot, but different
        # evidence assembly signatures (order + retrieval params).
        pricing_flags_a = {"restockEligible": True, "revision": 1}
        pricing_flags_b = {"restockEligible": True, "revision": 1}
        context_order_a = ["policy_version", "order_metadata", "pricing_flags", "tool_args"]
        context_order_b = ["policy_version", "pricing_flags", "order_metadata", "tool_args"]
        retrieval_params_a = {"retriever": "default", "reranker": "v1", "top_k": 4}
        retrieval_params_b = {"retriever": "default", "reranker": "v2", "top_k": 4}
    else:
        # B1 demo: dependency mutation (pricing flags change).
        pricing_flags_a = {"restockEligible": True, "revision": 1}
        pricing_flags_b = {"restockEligible": False, "revision": 2}
        context_order_a = ["policy_version", "order_metadata", "pricing_flags", "tool_args"]
        context_order_b = list(context_order_a)
        retrieval_params_a = {"retriever": "default", "reranker": "v1", "top_k": 4}
        retrieval_params_b = dict(retrieval_params_a)

    state_path = _state_file()
    _seed_state(state_path, customer_id=customer_id, payment_id=payment_id, amount_cents=amount_cents, reason=reason)
    executor_tools._STATE_FILE = str(state_path)

    print("=== BEFORE (Kybernis OFF) ===", flush=True)
    before = _run_before_without_kybernis(
        tool_args=tool_args,
        pricing_flags_a=pricing_flags_a,
        pricing_flags_b=pricing_flags_b,
    )
    print(json.dumps(before, indent=2), flush=True)

    print("\n=== AFTER (Kybernis ON) ===", flush=True)
    client = KybernisClient()
    if not client.base_url or not client.api_token or not client.tenant_id:
        print("Missing KYB_API_BASE / KYB_API_KEY / KYB_TENANT_ID for AFTER scenario.", flush=True)
        print("Video status: NOT_READY", flush=True)
        return 2

    try:
        run_id = client.get_or_create_run(
            input_context_ref=f"refund-video:{uuid.uuid4().hex}",
            workflow_version_id=workflow_version_id,
        )
    except Exception as exc:
        print(f"Could not create run: {exc}", flush=True)
        print("Tip: set KYB_WORKFLOW_VERSION_ID if your tenant has multiple workflow versions.", flush=True)
        print("Video status: NOT_READY", flush=True)
        return 2

    after = _run_after_with_kybernis(
        client=client,
        run_id=run_id,
        idempotency_key=f"intent:refund:video:{customer_id}:{payment_id}:{amount_cents}:{scenario_nonce}",
        policy_version=policy_version,
        order_metadata=order_metadata,
        pricing_flags_a=pricing_flags_a,
        pricing_flags_b=pricing_flags_b,
        context_order_a=context_order_a,
        context_order_b=context_order_b,
        retrieval_params_a=retrieval_params_a,
        retrieval_params_b=retrieval_params_b,
        tool_args=tool_args,
        requested_mode=requested_mode,
    )
    print(json.dumps(after, indent=2), flush=True)

    before_second_executed = bool(before.get("attempt_2_executed"))
    after_second_executed = bool(after.get("attempt_2_executed"))
    after_error = str(after.get("error") or "").strip()
    guard_2 = after.get("guard_2") if isinstance(after.get("guard_2"), dict) else {}
    drift_fields = guard_2.get("drift_fields_changed") if isinstance(guard_2, dict) else []
    if not isinstance(drift_fields, list):
        drift_fields = []
    b2_fields = {"evidence_order_signature", "retrieval_params_signature", "authorization_state_vector_signature", "context_ordering_signature"}
    b2_detected = any(str(item) in b2_fields for item in drift_fields)

    ready = before_second_executed and not after_second_executed and not after_error
    if drift_kind == "b2":
        ready = ready and b2_detected

    if ready:
        print(
            f"\nVIDEO_READY: {drift_kind.upper()} demo baseline executes attempt #2; Kybernis blocks attempt #2.",
            flush=True,
        )
        return 0

    print("\nVideo status: NOT_READY", flush=True)
    print("Expected: BEFORE attempt_2_executed=true and AFTER attempt_2_executed=false", flush=True)
    if after_error:
        print(f"AFTER error: {after_error}", flush=True)
    if drift_kind == "b2" and not b2_detected:
        print(f"B2 expected one of {sorted(b2_fields)} in drift_fields_changed, got: {drift_fields}", flush=True)
    return 1


if __name__ == "__main__":
    raise SystemExit(main())
