from __future__ import annotations

"""
local_kv_store.py (Python port of go-sdk/kvstore/local_kv_store.go)
-------------------------------------------------------------------

This module provides a *work-in-progress* Python implementation of the Bitcoin
SV on-chain key–value store originally implemented in Go.  Only a **minimal**
prototype is supplied at the moment – it fulfils the public API so that the
rest of the Python SDK can compile/import, yet the heavy blockchain logic is
still to be implemented.

Missing functionality is enumerated at the bottom of the file and returned via
`get_unimplemented_features()` so that build scripts / documentation can query
it programmatically.
"""

import base64
import copy
import json
import os
import re
from dataclasses import dataclass, field
from threading import Lock
from typing import Any, Dict, List, Optional, Tuple

from bsv.network.woc_client import WOCClient
from bsv.transaction.pushdrop import PushDrop

from .interfaces import (
    ErrEmptyContext,
    ErrInvalidKey,
    ErrInvalidValue,
    ErrInvalidWallet,
    KVStoreConfig,
    KVStoreInterface,
)

# ---------------------------------------------------------------------------
# Helper types
# ---------------------------------------------------------------------------


@dataclass
class _StoredValue:
    value: str
    # In the full implementation the fields below will reference on-chain
    # artefacts.  They are included here so that the public API (return types)
    # remain stable while the backing logic is developed.
    outpoint: str = ""  # txid.vout string – placeholder for now


# ---------------------------------------------------------------------------
# LocalKVStore prototype
# ---------------------------------------------------------------------------


class LocalKVStore(KVStoreInterface):
    """A *local* (in-memory) key–value store that mimics the Go behaviour.

    The real implementation must:
    1. Leverage *WalletInterface* to create PushDrop outputs on-chain
    2. Support optional encryption via wallet.Encrypt / wallet.Decrypt
    3. Collapse multiple values for the same key into a single UTXO when `set`
       is called repeatedly
    4. Handle removal by creating spending transactions that consume all
       matching outputs

    None of the above is done yet – instead we keep data in-memory so that unit
    tests targeting higher-level components can progress.
    """

    _UNIMPLEMENTED: list[str] = [
        # BEEF / AtomicBEEF parsing is now implemented
        # Retention period & basket name support is now implemented
    ]

    # NOTE: We do *not* attempt to replicate the rich context propagation of Go
    # right now – the `ctx` parameter is accepted but not inspected.

    def __init__(self, config: KVStoreConfig):
        if config.wallet is None:
            raise ErrInvalidWallet("wallet cannot be None")
        if not config.context:
            raise ErrEmptyContext("context cannot be empty")

        self._wallet = config.wallet
        self._context = config.context
        self._retention_period: int = int(getattr(config, "retention_period", 0) or 0)
        self._basket_name: str = getattr(config, "basket_name", "") or self._context
        self._protocol = re.sub(r"[^A-Za-z0-9 ]", "", self._context).replace(" ", "")
        self._originator = config.originator
        self._encrypt = bool(config.encrypt)
        # TS/GO-style defaults
        self._default_fee_rate: int | None = getattr(config, "fee_rate", None)
        self._default_ca: dict | None = getattr(config, "default_ca", None)
        self._lock_position: str = getattr(config, "lock_position", "before") or "before"
        # Remove _use_local_store and _store except for test hooks
        self._lock = Lock()
        # Key-level locks (per-key serialization)
        self._key_locks: dict[str, Lock] = {}
        self._key_locks_guard: Lock = Lock()
        # Options
        self._accept_delayed_broadcast: bool = bool(
            getattr(config, "accept_delayed_broadcast", False) or getattr(config, "acceptDelayedBroadcast", False)
        )
        # Cache: recently created BEEF per key to avoid WOC on immediate get
        self._recent_beef_by_key: dict[str, tuple[list, bytes]] = {}

    # ---------------------------------------------------------------------
    # Helper methods
    # ---------------------------------------------------------------------

    def _get_protocol(self, key: str) -> dict:
        """Returns the wallet protocol for the given key (GO pattern).

        This method mirrors the Go SDK's getProtocol() implementation.
        It returns only the protocol structure, as keyID is always the same
        as the key parameter and should be passed separately.

        Args:
            key: The key string (not used in protocol generation, but kept for API consistency)

        Returns:
            dict: Protocol dict with 'securityLevel' and 'protocol' keys.
                  securityLevel is 2 (SecurityLevelEveryAppAndCounterparty).
                  protocol is derived from the context.

        Note:
            keyID is not included in the return value as it's always the same
            as the key parameter. This follows the Go SDK pattern.
        """
        return {"securityLevel": 2, "protocol": self._protocol}

    # ---------------------------------------------------------------------
    # Public API
    # ---------------------------------------------------------------------

    def get(self, ctx: Any, key: str, default_value: str = "") -> str:
        if not key:
            raise ErrInvalidKey(KEY_EMPTY_MSG)
        self._acquire_key_lock(key)
        try:
            value = self._get_onchain_value(ctx, key)
            if value is not None:
                return value
            return default_value
        finally:
            self._release_key_lock(key)

    def _get_onchain_value(self, ctx: Any, key: str) -> str | None:
        """Retrieve value from on-chain outputs (BEEF/PushDrop)."""
        outputs, beef_bytes = self._lookup_outputs_for_get(ctx, key)
        if not outputs:
            return None

        most_recent = outputs[-1]
        locking_script = self._extract_locking_script_from_output(beef_bytes, most_recent)
        if not locking_script:
            return None

        decoded = PushDrop.decode(locking_script)
        if not decoded or not isinstance(decoded.get("fields"), list) or not decoded["fields"]:
            return None

        first_field = decoded["fields"][0]
        return self._process_decoded_field(ctx, first_field)

    def _process_decoded_field(self, ctx: Any, first_field: Any) -> str | None:
        """Process a decoded field, handling encryption if enabled."""
        if self._should_return_encrypted(first_field):
            return self._format_encrypted_field(first_field)

        if self._should_attempt_decryption():
            decrypted = self._attempt_decryption(ctx, first_field)
            if decrypted is not None:
                return decrypted

        return self._decode_plaintext_field(first_field)

    def _should_return_encrypted(self, first_field: Any) -> bool:
        """Check if field should be returned in encrypted form."""
        return (
            self._encrypt
            and isinstance(self._default_ca, dict)
            and self._default_ca
            and (
                isinstance(first_field, (bytes, bytearray))
                or (isinstance(first_field, str) and not first_field.startswith("enc:"))
            )
        )

    def _format_encrypted_field(self, first_field: Any) -> str | None:
        """Format field as encrypted string."""
        try:
            if isinstance(first_field, (bytes, bytearray)):
                return "enc:" + base64.b64encode(bytes(first_field)).decode("ascii")
            elif isinstance(first_field, str):
                if first_field.startswith("enc:"):
                    return first_field
                return "enc:" + base64.b64encode(first_field.encode("utf-8")).decode("ascii")
        except Exception:
            pass
        return None

    def _should_attempt_decryption(self) -> bool:
        """Check if decryption should be attempted."""
        return self._encrypt and isinstance(self._default_ca, dict) and self._default_ca

    def _attempt_decryption(self, ctx: Any, first_field: Any) -> str | None:
        """Attempt to decrypt the field."""
        try:
            ciphertext = self._normalize_ciphertext(first_field)
            if not ciphertext:
                return None

            ca_args = self._merge_default_ca(None)
            protocol_id, key_id, counterparty = self._extract_encryption_params_from_ca(ca_args)

            dec_res = (
                self._wallet.decrypt(
                    ctx,
                    {
                        "encryption_args": {
                            "protocolID": protocol_id,
                            "keyID": key_id,
                            "counterparty": counterparty,
                        },
                        "ciphertext": ciphertext,
                    },
                    self._originator,
                )
                or {}
            )

            pt = dec_res.get("plaintext")
            if isinstance(pt, (bytes, bytearray)):
                return pt.decode("utf-8")
        except Exception:
            pass
        return None

    def _normalize_ciphertext(self, first_field: Any) -> bytes:
        """Normalize field to ciphertext bytes."""
        if isinstance(first_field, (bytes, bytearray)):
            first_field_bytes = bytes(first_field)
            if first_field_bytes.startswith(b"enc:"):
                return base64.b64decode(first_field_bytes[4:])
            else:
                return first_field_bytes
        elif isinstance(first_field, str):
            if first_field.startswith("enc:"):
                return base64.b64decode(first_field[4:])
            else:
                return first_field.encode("utf-8")
        return b""

    def _extract_encryption_params_from_ca(self, ca_args: dict) -> tuple:
        """Extract encryption parameters from CA args."""
        pd_opts = ca_args.get("pushdrop") or {}
        protocol_id = (
            ca_args.get("protocol_id")
            or ca_args.get("protocolID")
            or pd_opts.get("protocol_id")
            or pd_opts.get("protocolID")
        )
        key_id = ca_args.get("key_id") or ca_args.get("keyID") or pd_opts.get("key_id") or pd_opts.get("keyID")
        counterparty = ca_args.get("counterparty") or pd_opts.get("counterparty") or {"type": 2}  # Default to SELF (2)
        return protocol_id, key_id, counterparty

    def _decode_plaintext_field(self, first_field: Any) -> str | None:
        """Decode field as plaintext."""
        try:
            if isinstance(first_field, (bytes, bytearray)):
                return first_field.decode("utf-8")
            elif isinstance(first_field, str):
                return first_field
        except Exception:
            pass
        return None

    def _lookup_outputs_for_get(self, _ctx: Any, key: str) -> tuple[list, bytes]:
        """Lookup outputs for get operation using multiple fallback strategies."""
        # Fast-path: return locally cached BEEF right after set
        cached_result = self._get_cached_outputs(key)
        if cached_result:
            return cached_result

        # Primary lookup via wallet
        outputs, beef_bytes = self._lookup_via_wallet(key)

        # Fallback 1: build BEEF from WOC if no BEEF but have outputs
        if not beef_bytes and outputs:
            beef_bytes = self._build_beef_from_woc_outputs(outputs)

        # Fallback 2: scan WOC address histories for PushDrop outputs
        if not outputs and not beef_bytes:
            outputs, beef_bytes = self._scan_woc_for_pushdrop_outputs()

        return outputs, beef_bytes

    def _get_cached_outputs(self, key: str) -> tuple[list, bytes] | None:
        """Return cached outputs if available."""
        cached = self._recent_beef_by_key.get(key)
        if cached:
            outs, beef = cached
            if outs and beef:
                return outs, beef
        return None

    def _lookup_via_wallet(self, key: str) -> tuple[list, bytes]:
        """Primary lookup using wallet.list_outputs."""
        args = self._build_list_outputs_args(key)
        lo = self._wallet.list_outputs(args, self._originator) or {}
        outputs = lo.get("outputs") or []
        beef_bytes = lo.get("BEEF") or b""
        return outputs, beef_bytes

    def _build_list_outputs_args(self, key: str) -> dict:
        """Build arguments for wallet.list_outputs call."""
        args = {
            "basket": self._context,
            "tags": [key],
            "include": ENTIRE_TXS,
            "limit": 10,
        }

        # Forward derivation defaults for derived-address lookup
        try:
            ca_args = self._merge_default_ca(None)
            pd_opts = ca_args.get("pushdrop") or {}
            prot = self._extract_protocol_id(ca_args, pd_opts)
            kid = self._extract_key_id(ca_args, pd_opts)
            cpty = ca_args.get("counterparty") or pd_opts.get("counterparty")

            if prot is not None:
                args["protocolID"] = prot
            if kid is not None:
                args["keyID"] = kid
            if cpty is not None:
                args["counterparty"] = cpty
        except Exception:
            pass

        return args

    def _extract_protocol_id(self, ca_args: dict, pd_opts: dict) -> Any:
        """Extract protocol ID from CA args."""
        return (
            ca_args.get("protocol_id")
            or ca_args.get("protocolID")
            or pd_opts.get("protocol_id")
            or pd_opts.get("protocolID")
        )

    def _extract_key_id(self, ca_args: dict, pd_opts: dict) -> Any:
        """Extract key ID from CA args."""
        return ca_args.get("key_id") or ca_args.get("keyID") or pd_opts.get("key_id") or pd_opts.get("keyID")

    def _build_beef_from_woc_outputs(self, outputs: list) -> bytes:
        """Build BEEF from WOC outputs as fallback."""
        try:
            timeout = int(os.getenv("WOC_TIMEOUT", "10"))
            return self._build_beef_v2_from_woc_outputs(outputs, timeout=timeout)
        except Exception:
            return b""

    def _scan_woc_for_pushdrop_outputs(self) -> tuple[list, bytes]:
        """Scan WOC for PushDrop outputs as final fallback."""
        try:
            candidates = self._get_address_candidates()
            woc_api = os.environ.get("WOC_API_KEY") or ""
            headers = {"Authorization": woc_api, "woc-api-key": woc_api} if woc_api else {}
            timeout = int(os.getenv("WOC_TIMEOUT", "10"))

            matched_outputs: list[dict] = []
            matched_tx_hexes: list[str] = []
            seen_txids: set = set()

            for addr, pub_hex in candidates:
                self._scan_address_for_pushdrop_outputs(
                    addr, pub_hex, headers, timeout, seen_txids, matched_outputs, matched_tx_hexes
                )

            if matched_outputs and matched_tx_hexes:
                unique_tx_hexes = list(dict.fromkeys(matched_tx_hexes))
                from bsv.beef import build_beef_v2_from_raw_hexes

                beef_bytes = build_beef_v2_from_raw_hexes(unique_tx_hexes)
                return matched_outputs, beef_bytes

        except Exception as e_fallback2:
            print(f"[KV WOC] fallback-2 scan failed: {e_fallback2}")

        return [], b""

    def _create_protocol_object(self, ca_args: dict, pd_opts: dict) -> Any:
        """Create protocol object from CA args and pushdrop options."""
        prot = self._extract_protocol_id(ca_args, pd_opts)
        if prot is None:
            return None

        from bsv.wallet.key_deriver import Protocol

        return Protocol(prot.get("securityLevel", 0), prot.get("protocol", "")) if isinstance(prot, dict) else prot

    def _derive_address_components(self, protocol_obj: Any, kid: Any, cpty: Any) -> tuple[Any, str | None, str | None]:
        """Derive address components from protocol, key ID, and counterparty."""
        if protocol_obj is None or kid is None:
            return None, None, None

        cp_norm = (
            self._wallet._normalize_counterparty(cpty) if hasattr(self._wallet, "_normalize_counterparty") else None
        )

        try:
            derived_pub = self._wallet.key_deriver.derive_public_key(protocol_obj, kid, cp_norm, for_self=False)
            derived_addr = derived_pub.address()
            derived_pub_hex = derived_pub.hex()
            return derived_pub, derived_addr, derived_pub_hex
        except Exception:
            return None, None, None

    def _get_master_address(self) -> str | None:
        """Get the master address from the wallet."""
        try:
            return self._wallet.public_key.address()
        except Exception:
            return None

    def _build_address_candidates(
        self, master_addr: str | None, derived_addr: str | None, derived_pub_hex: str | None
    ) -> list[tuple[str, str | None]]:
        """Build the list of address candidates."""
        candidates: list[tuple[str, str | None]] = []

        if master_addr:
            candidates.append((master_addr, derived_pub_hex))

        # Optional: if LocalKVStore.context is an address distinct from above, include it
        try:
            basket_ctx = self._context
            if isinstance(basket_ctx, str) and basket_ctx:
                is_new = (basket_ctx != master_addr) and (basket_ctx != derived_addr)
                if is_new and self._looks_like_address(basket_ctx):
                    candidates.append((basket_ctx, derived_pub_hex))
        except Exception:
            pass

        if derived_addr:
            candidates.append((derived_addr, derived_pub_hex))

        return candidates

    def _get_address_candidates(self) -> list[tuple[str, str | None]]:
        """Get list of addresses to scan for PushDrop outputs."""
        ca_args = self._merge_default_ca(None)
        pd_opts = ca_args.get("pushdrop") or {}
        cpty = ca_args.get("counterparty") or pd_opts.get("counterparty")

        protocol_obj = self._create_protocol_object(ca_args, pd_opts)
        kid = self._extract_key_id(ca_args, pd_opts)
        _, derived_addr, derived_pub_hex = self._derive_address_components(protocol_obj, kid, cpty)
        master_addr = self._get_master_address()

        return self._build_address_candidates(master_addr, derived_addr, derived_pub_hex)

    def _scan_address_for_pushdrop_outputs(
        self,
        addr: str,
        pub_hex: str,
        headers: dict,
        timeout: int,
        seen_txids: set,
        matched_outputs: list,
        matched_tx_hexes: list,
    ) -> None:
        """Scan a WOC address for PushDrop outputs matching the given public key."""
        try:
            txs = self._fetch_address_history(addr, headers, timeout)
            if txs is None:
                return

            txids = self._extract_txids_from_history(txs)
            for txid in [x for x in txids if x][:50]:
                if txid in seen_txids:
                    continue
                seen_txids.add(txid)

                rawtx = self._fetch_raw_transaction(txid, headers, timeout)
                if not rawtx:
                    continue

                self._process_transaction_for_pushdrop(txid, rawtx, pub_hex, addr, matched_outputs, matched_tx_hexes)
        except Exception as e_addr_loop:
            print(f"[KV WOC] address loop error for {addr}: {e_addr_loop}")

    def _fetch_address_history(self, addr: str, headers: dict, timeout: int):
        """Fetch transaction history for an address from WOC."""
        import requests

        base = f"https://api.whatsonchain.com/v1/bsv/main/address/{addr}"
        hist_endpoints = [
            f"{base}/confirmed/history",
            f"{base}/history",
        ] + [f"{base}/txs/{p}" for p in range(3)]

        for hist_url in hist_endpoints:
            try:
                print(f"[KV WOC] try history endpoint: {hist_url}")
                r = requests.get(hist_url, headers=headers, timeout=timeout)
                if r.status_code == 404:
                    continue
                r.raise_for_status()
                resp = r.json() or []
                txs = self._normalize_history_response(resp)
                if txs is not None:
                    return txs
            except Exception:
                continue

        # Fallback to UTXO endpoint
        return self._fetch_address_utxos(base, headers, timeout)

    def _normalize_history_response(self, resp):
        """Normalize various WOC history response shapes."""
        if isinstance(resp, dict):
            for key in ["result", "transactions", "txs", "history"]:
                if isinstance(resp.get(key), list):
                    return resp[key]
            return []
        return resp

    def _fetch_address_utxos(self, base_url: str, headers: dict, timeout: int):
        """Fetch UTXOs as a fallback for transaction history."""
        import requests

        utxo_url = f"{base_url}/unspent"
        try:
            print(f"[KV WOC] fallback to UTXO endpoint: {utxo_url}")
            r = requests.get(utxo_url, headers=headers, timeout=timeout)
            r.raise_for_status()
            return r.json() or []
        except Exception as e:
            print(f"[KV WOC] UTXO fetch failed: {e}")
            return None

    def _extract_txids_from_history(self, txs: list) -> list:
        """Extract transaction IDs from history response."""
        txids = []
        for t in txs:
            if isinstance(t, str) and len(t) == 64:
                txids.append(t)
            elif isinstance(t, dict):
                txids.append(t.get("tx_hash") or t.get("txid") or t.get("hash") or "")
        return txids

    def _fetch_raw_transaction(self, txid: str, headers: dict, timeout: int):
        """Fetch raw transaction hex from WOC."""
        import requests

        raw_candidates = [
            f"https://api.whatsonchain.com/v1/bsv/main/tx/{txid}/hex",
            f"https://api.whatsonchain.com/v1/bsv/main/tx/{txid}",
            f"https://api.whatsonchain.com/v1/bsv/main/tx/raw/{txid}",
        ]

        for raw_url in raw_candidates:
            try:
                print(f"[KV WOC] try tx endpoint: {raw_url}")
                rr = requests.get(raw_url, headers=headers, timeout=timeout)
                if rr.status_code == 404:
                    continue
                rr.raise_for_status()

                ctype = rr.headers.get("Content-Type", "")
                if "application/json" in ctype:
                    jd = rr.json() or {}
                    rawtx = jd.get("hex") or jd.get("rawtx") or jd.get("data")
                else:
                    rawtx = rr.text.strip()

                if isinstance(rawtx, str) and len(rawtx) >= 2:
                    return rawtx
            except Exception:
                continue

        print(f"[KV WOC] raw fetch failed for {txid}")
        return None

    def _process_transaction_for_pushdrop(
        self, txid: str, rawtx: str, pub_hex: str, addr: str, matched_outputs: list, matched_tx_hexes: list
    ) -> None:
        """Process a transaction to find PushDrop outputs for the given public key."""
        from bsv.transaction import Transaction
        from bsv.utils import Reader

        try:
            tx = Transaction.from_reader(Reader(bytes.fromhex(rawtx)))
        except Exception as e:
            print(f"[KV WOC] tx parse failed for {txid}: {e}")
            return

        for vout_idx, out in enumerate(tx.outputs):
            try:
                ls_bytes = out.locking_script.to_bytes()
                if self._is_pushdrop_for_pub(ls_bytes, pub_hex):
                    matched_outputs.append(
                        {
                            "outputIndex": vout_idx,
                            "satoshis": out.satoshis,
                            "lockingScript": ls_bytes.hex(),
                            "spendable": True,
                            "outputDescription": "WOC scan (PushDrop)",
                            "basket": addr,
                            "tags": [],
                            "customInstructions": None,
                            "txid": tx.txid(),
                        }
                    )
                    matched_tx_hexes.append(rawtx)
                    break
            except Exception as e:
                print(f"[KV WOC] vout scan error in {txid}@{vout_idx}: {e}")

    def _looks_like_address(self, addr: str) -> bool:
        """Best-effort check if a string is a Base58Check address (no network assert)."""
        try:
            if not isinstance(addr, str) or len(addr) < 26 or len(addr) > 50:
                return False
            from bsv.utils.base58_utils import from_base58_check

            _ = from_base58_check(addr)
            return True
        except Exception:
            return False

    def _extract_locking_script_from_output(self, beef_bytes: bytes, output: dict) -> bytes:
        locking_script = output.get("lockingScript") or b""
        if not beef_bytes:
            return locking_script
        try:
            from bsv.transaction import parse_beef_ex

            beef, subject, last_tx = parse_beef_ex(beef_bytes)
            txid_hint = output.get("txid")
            match_tx = self._find_tx_by_subject(beef, subject)
            if match_tx is not None:
                vout = int(output.get("outputIndex", 0))
                if 0 <= vout < len(match_tx.outputs):
                    return match_tx.outputs[vout].locking_script.to_bytes()  # Scriptオブジェクトからbytesを取得
            match_tx = self._find_tx_by_txid_hint(beef, txid_hint)
            if match_tx is not None:
                vout = int(output.get("outputIndex", 0))
                if 0 <= vout < len(match_tx.outputs):
                    return match_tx.outputs[vout].locking_script.to_bytes()  # Scriptオブジェクトからbytesを取得
            if last_tx is not None:
                vout = int(output.get("outputIndex", 0))
                if 0 <= vout < len(last_tx.outputs):
                    return last_tx.outputs[vout].locking_script.to_bytes()  # Scriptオブジェクトからbytesを取得
        except Exception:
            pass
        return locking_script

    def _find_tx_by_subject(self, beef, subject):
        if not subject:
            return None
        btxs = beef.find_transaction(subject)
        if btxs and getattr(btxs, "tx_obj", None):
            return btxs.tx_obj
        return None

    def _find_tx_by_txid_hint(self, beef, txid_hint):
        if not (txid_hint and isinstance(txid_hint, str)):
            return None
        btx = beef.find_transaction(txid_hint)
        if btx and getattr(btx, "tx_obj", None):
            return btx.tx_obj
        return None

    def set(self, ctx: Any, key: str, value: str, ca_args: dict = None) -> str:
        if not key:
            raise ErrInvalidKey(KEY_EMPTY_MSG)
        if not value:
            raise ErrInvalidValue("Value cannot be empty")

        self._acquire_key_lock(key)
        try:
            return self._execute_set_operation(ctx, key, value, ca_args)
        finally:
            self._release_key_lock(key)

    def _execute_set_operation(self, ctx: Any, key: str, value: str, ca_args: dict) -> str:
        """Execute the set operation with all required steps."""
        ca_args = self._merge_default_ca(ca_args)
        print(f"[TRACE] [set] ca_args: {ca_args}")

        # Prepare transaction components
        outs, input_beef = self._lookup_outputs_for_set(ctx, key, ca_args)
        locking_script = self._build_locking_script(ctx, key, value, ca_args)
        inputs_meta = self._prepare_inputs_meta(key, outs, ca_args)
        print(f"[TRACE] [set] inputs_meta after _prepare_inputs_meta: {inputs_meta}")

        # Create and sign transaction
        create_args = self._build_create_action_args_set(key, value, locking_script, inputs_meta, input_beef, ca_args)
        create_args["inputs"] = inputs_meta
        if ca_args and "use_woc" in ca_args:
            create_args["use_woc"] = ca_args["use_woc"]

        ca = self._wallet.create_action(create_args, self._originator) or {}
        signable = (ca.get("signableTransaction") or {}) if isinstance(ca, dict) else {}
        signable_tx_bytes = signable.get("tx") or b""

        signed_tx_bytes = None
        if inputs_meta:
            signed_tx_bytes = self._sign_and_relinquish_set(
                ctx, key, outs, inputs_meta, signable, signable_tx_bytes, input_beef
            )

        # Cache BEEF for immediate retrieval
        tx_bytes = signed_tx_bytes or signable_tx_bytes
        self._build_and_cache_beef(key, locking_script, tx_bytes)

        # Broadcast and return result
        self._wallet.internalize_action({"tx": tx_bytes}, self._originator)
        return self._extract_txid_from_bytes(tx_bytes, key)

    def _build_and_cache_beef(self, key: str, locking_script: bytes, tx_bytes: bytes) -> None:
        """Build BEEF from transaction and cache it for immediate retrieval."""
        try:
            import binascii

            from bsv.beef import build_beef_v2_from_raw_hexes
            from bsv.script.script import Script
            from bsv.transaction import Transaction, TransactionOutput
            from bsv.utils import Reader

            tx, tx_hex = self._parse_or_create_transaction(tx_bytes, locking_script)
            beef_now = build_beef_v2_from_raw_hexes([tx_hex]) if tx_hex else b""

            if beef_now:
                locking_script_hex = (
                    locking_script.hex() if isinstance(locking_script, (bytes, bytearray)) else str(locking_script)
                )
                recent_outs = [
                    {
                        "outputIndex": 0,
                        "satoshis": 1,
                        "lockingScript": locking_script_hex,
                        "spendable": True,
                        "outputDescription": "KV set (local)",
                        "basket": self._context,
                        "tags": [key, "kv", "set"],
                        "customInstructions": None,
                        "txid": tx.txid() if hasattr(tx, "txid") else "",
                    }
                ]
                self._recent_beef_by_key[key] = (recent_outs, beef_now)
        except Exception as e_beef:
            print(f"[KV set] build immediate BEEF failed: {e_beef}")

    def _parse_or_create_transaction(self, tx_bytes: bytes, locking_script: bytes):
        """Parse transaction from bytes or create a minimal transaction."""
        import binascii

        from bsv.script.script import Script
        from bsv.transaction import Transaction, TransactionOutput
        from bsv.utils import Reader

        if tx_bytes:
            try:
                tx = Transaction.from_reader(Reader(tx_bytes))
                tx_hex = binascii.hexlify(tx_bytes).decode()
                return tx, tx_hex
            except Exception:
                pass

        # Fallback: synthesize a minimal transaction
        try:
            ls_bytes = (
                locking_script if isinstance(locking_script, (bytes, bytearray)) else bytes.fromhex(str(locking_script))
            )
        except Exception:
            ls_bytes = b""

        tx = Transaction()
        tx.outputs = [TransactionOutput(Script(ls_bytes), 1)]
        tx_hex = tx.serialize().hex()
        return tx, tx_hex

    def _extract_txid_from_bytes(self, tx_bytes: bytes, key: str) -> str:
        """Extract txid from transaction bytes or return fallback."""
        try:
            from bsv.transaction import Transaction
            from bsv.utils import Reader

            if tx_bytes:
                tx = Transaction.from_reader(Reader(tx_bytes))
                return f"{tx.txid()}.0"
        except Exception:
            pass
        return f"{key}.0"

    def _build_locking_script(self, _ctx: Any, key: str, value: str, ca_args: dict = None) -> str:
        ca_args = self._merge_default_ca(ca_args)
        field_bytes = self._encrypt_value_if_needed(key, value, ca_args)

        protocol_id, key_id, counterparty = self._extract_pushdrop_params(ca_args)

        pd = PushDrop(self._wallet, self._originator)
        return pd.lock(
            [field_bytes],
            protocol_id,
            key_id,
            counterparty,
            for_self=True,
            include_signature=True,
            lock_position="before",
        )

    def _encrypt_value_if_needed(self, key: str, value: str, ca_args: dict) -> bytes:
        """Encrypt value if encryption is enabled, otherwise return plaintext bytes."""
        if not self._encrypt:
            return value.encode("utf-8")

        protocol_id, key_id, counterparty = self._extract_encryption_params(key, ca_args)
        if not (protocol_id and key_id):
            return value.encode("utf-8")

        return self._perform_encryption(value, protocol_id, key_id, counterparty)

    def _extract_encryption_params(self, key: str, ca_args: dict) -> tuple:
        """Extract protocol_id, key_id, and counterparty for encryption."""
        protocol_id = ca_args.get("protocol_id") or ca_args.get("protocolID") or self._get_protocol(key)
        key_id = ca_args.get("key_id") or ca_args.get("keyID") or key
        counterparty = ca_args.get("counterparty") or {"type": 2}  # Default to SELF (2)
        return protocol_id, key_id, counterparty

    def _perform_encryption(self, value: str, protocol_id: Any, key_id: Any, counterparty: dict) -> bytes:
        """Perform encryption and return ciphertext bytes."""
        is_self = isinstance(counterparty, dict) and counterparty.get("type") == 0
        encrypt_args = {
            "encryption_args": {
                "protocolID": protocol_id,
                "keyID": key_id,
                "counterparty": counterparty,
                "forSelf": is_self,
            },
            "plaintext": value.encode("utf-8"),
        }
        encrypt_result = self._wallet.encrypt(encrypt_args, self._originator)

        if "ciphertext" in encrypt_result:
            ciphertext = encrypt_result["ciphertext"]
            if isinstance(ciphertext, (list, bytes, bytearray)):
                return bytes(ciphertext)

        # Fallback to plaintext if encryption fails
        return value.encode("utf-8")

    def _extract_pushdrop_params(self, ca_args: dict) -> tuple:
        """Extract protocol_id, key_id, and counterparty for PushDrop."""
        pd_opts = ca_args.get("pushdrop") or {}
        protocol_id = (
            ca_args.get("protocol_id")
            or ca_args.get("protocolID")
            or pd_opts.get("protocol_id")
            or pd_opts.get("protocolID")
        )
        key_id = ca_args.get("key_id") or ca_args.get("keyID") or pd_opts.get("key_id") or pd_opts.get("keyID")
        counterparty = ca_args.get("counterparty", pd_opts.get("counterparty"))
        return protocol_id, key_id, counterparty

    def _lookup_outputs_for_set(self, _ctx: Any, key: str, ca_args: dict | None = None) -> tuple[list, bytes]:
        ca_args = self._merge_default_ca(ca_args)
        address = self._context
        # Preserve original behaviour (basket/tags) and pass-through ca_args for optional derived lookup
        args = {
            "basket": address,
            "tags": [key],
            "include": ENTIRE_TXS,
            "limit": 100,
        }
        # Non-intrusive: forward protocolID/keyID/counterparty only if present
        pd_opts = ca_args.get("pushdrop") or {}
        prot = (
            ca_args.get("protocol_id")
            or ca_args.get("protocolID")
            or pd_opts.get("protocol_id")
            or pd_opts.get("protocolID")
        )
        kid = ca_args.get("key_id") or ca_args.get("keyID") or pd_opts.get("key_id") or pd_opts.get("keyID")
        cpty = ca_args.get("counterparty") or pd_opts.get("counterparty")
        if prot is not None:
            args["protocol_id"] = prot
        if kid is not None:
            args["key_id"] = kid
        if cpty is not None:
            args["counterparty"] = cpty
        lo = self._wallet.list_outputs(args, self._originator) or {}
        outs = [o for o in lo.get("outputs") or [] if not o.get("error")]
        input_beef = lo.get("BEEF") or b""
        if not input_beef and outs:
            try:
                timeout = int(os.getenv("WOC_TIMEOUT", "10"))
                input_beef = self._build_beef_v2_from_woc_outputs(outs, timeout=timeout)
            except Exception:
                input_beef = b""
        return outs, input_beef

    def _build_create_action_args_set(
        self, key: str, value: str, locking_script: bytes, inputs_meta: list, input_beef: bytes, ca_args: dict = None
    ) -> dict:
        ca_args = self._merge_default_ca(ca_args)
        pd_opts = ca_args.get("pushdrop") or {}
        protocol_id = (
            ca_args.get("protocol_id")
            or ca_args.get("protocolID")
            or pd_opts.get("protocol_id")
            or pd_opts.get("protocolID")
        )
        key_id = ca_args.get("key_id") or ca_args.get("keyID") or pd_opts.get("key_id") or pd_opts.get("keyID")
        counterparty = ca_args.get("counterparty", pd_opts.get("counterparty"))
        fee_rate = ca_args.get("feeRate", ca_args.get("fee_rate", self._default_fee_rate))
        fields = [value.encode("utf-8")]
        # locking_script: always hex string for Go/TS parity
        if isinstance(locking_script, bytes):
            locking_script_hex = locking_script.hex()
        else:
            locking_script_hex = locking_script
        return {
            "labels": ["kv", "set"],
            "pushdrop": {
                "fields": fields,
                # Provide both snake_case and camelCase for compatibility
                "protocol_id": protocol_id,
                "protocolID": protocol_id,
                "key_id": key_id,
                "keyID": key_id,
                "counterparty": counterparty,
                "forSelf": True,
                "include_signature": True,  # Restored: Enable PushDrop signature for normal operation
                "lock_position": "before",
            },
            "inputs_meta": inputs_meta,
            "input_beef": input_beef,
            "outputs": [
                {
                    "lockingScript": locking_script_hex,
                    "satoshis": 1,
                    "tags": [key, "kv", "set"],
                    "basket": self._context,
                    "outputDescription": (
                        {"retentionSeconds": self._retention_period} if int(self._retention_period or 0) > 0 else ""
                    ),
                }
            ],
            "feeRate": fee_rate,
            "options": {
                "acceptDelayedBroadcast": self._accept_delayed_broadcast,
                "randomizeOutputs": False,
            },
        }

    def _sign_and_relinquish_set(
        self,
        _ctx: Any,
        key: str,
        outs: list,
        inputs_meta: list,
        signable: dict,
        signable_tx_bytes: bytes,
        input_beef: bytes,
    ) -> bytes | None:
        spends = self._prepare_spends(key, inputs_meta, signable_tx_bytes, input_beef)
        try:
            spends_str_keys = {str(int(k)): v for k, v in (spends or {}).items()}
            res = self._wallet.sign_action(
                {
                    "spends": spends_str_keys,
                    "reference": signable.get("reference") or b"",
                    "tx": signable_tx_bytes,
                },
                self._originator,
            )
            return (res or {}).get("tx") if isinstance(res, dict) else None
        except Exception:
            for o in outs:
                try:
                    self._wallet.relinquish_output(
                        {
                            "basket": self._context,
                            "output": {
                                "txid": (
                                    bytes.fromhex(o.get("txid", "00" * 32))
                                    if isinstance(o.get("txid"), str)
                                    else (o.get("txid") or b"\x00" * 32)
                                ),
                                "index": int(o.get("outputIndex", 0)),
                            },
                        },
                        self._originator,
                    )
                except Exception:
                    pass
            return None

    def remove(self, ctx: Any, key: str) -> list[str]:  # NOSONAR - Complexity (17), requires refactoring
        if not key:
            raise ErrInvalidKey(KEY_EMPTY_MSG)
        self._acquire_key_lock(key)
        removed: list[str] = []
        loop_guard = 0
        last_count = None
        try:
            while True:
                if loop_guard > 10:
                    break
                loop_guard += 1
                outs, input_beef, total_outputs = self._lookup_outputs_for_remove(ctx, key)
                count = len(outs)
                if count == 0:
                    break
                if last_count is not None and count >= last_count:
                    break
                last_count = count
                inputs_meta = self._prepare_inputs_meta(key, outs)
                txid = self._onchain_remove_flow(ctx, key, inputs_meta, input_beef)
                if isinstance(txid, str) and txid:
                    removed.append(txid)
                # TS parity: break when outputs processed equals totalOutputs
                try:
                    if isinstance(total_outputs, int) and count == total_outputs:
                        break
                except Exception:
                    pass
            return removed
        finally:
            self._release_key_lock(key)

    def _lookup_outputs_for_remove(self, _ctx: Any, key: str) -> tuple[list, bytes, int | None]:
        lo = (
            self._wallet.list_outputs(
                {
                    "basket": self._context,
                    "tags": [key],
                    "include": ENTIRE_TXS,
                    "limit": 100,
                },
                self._originator,
            )
            or {}
        )
        outs = lo.get("outputs") or []
        input_beef = lo.get("BEEF") or b""
        total_outputs = None
        try:
            total_outputs = lo.get("totalOutputs") or lo.get("total_outputs")
            if isinstance(total_outputs, str) and total_outputs.isdigit():
                total_outputs = int(total_outputs)
        except Exception:
            total_outputs = None
        if not input_beef and outs:
            try:
                timeout = int(os.getenv("WOC_TIMEOUT", "10"))
                input_beef = self._build_beef_v2_from_woc_outputs(outs, timeout=timeout)
            except Exception:
                input_beef = b""
        return outs, input_beef, total_outputs

    def _onchain_remove_flow(self, _ctx: Any, key: str, inputs_meta: list, input_beef: bytes) -> str | None:
        ca_res = (
            self._wallet.create_action(
                {
                    "labels": ["kv", "remove"],
                    "description": f"kvstore remove {key}",
                    "inputs": inputs_meta,
                    "inputBEEF": input_beef,
                    "options": {"acceptDelayedBroadcast": self._accept_delayed_broadcast},
                },
                self._originator,
            )
            or {}
        )
        signable = (ca_res.get("signableTransaction") or {}) if isinstance(ca_res, dict) else {}
        signable_tx_bytes = signable.get("tx") or b""
        reference = signable.get("reference") or b""
        spends = self._prepare_spends(key, inputs_meta, signable_tx_bytes, input_beef)
        spends_str = {str(int(k)): v for k, v in (spends or {}).items()}
        res = self._wallet.sign_action({"spends": spends_str, "reference": reference}, self._originator) or {}
        signed_tx_bytes = res.get("tx") if isinstance(res, dict) else None
        internalize_result = self._wallet.internalize_action(
            {"tx": signed_tx_bytes or signable_tx_bytes}, self._originator
        )
        parsed_txid = None
        try:
            from bsv.transaction import Transaction
            from bsv.utils import Reader

            tx_bytes_final = signed_tx_bytes or signable_tx_bytes
            if tx_bytes_final:
                t = Transaction.from_reader(Reader(tx_bytes_final))
                parsed_txid = t.txid()
        except Exception:
            pass
        # Use parsed txid if available, otherwise use txid from internalize_action (for mocks)
        if parsed_txid:
            return parsed_txid
        if isinstance(internalize_result, dict) and internalize_result.get("txid"):
            return internalize_result["txid"]
        return None

    # ------------------------------
    # Key-level locking helpers
    # ------------------------------
    def _acquire_key_lock(self, key: str) -> None:
        try:
            with self._key_locks_guard:
                lk = self._key_locks.get(key)
                if lk is None:
                    lk = Lock()
                    self._key_locks[key] = lk
            lk.acquire()
        except Exception:
            pass

    def _release_key_lock(self, key: str) -> None:
        try:
            lk = self._key_locks.get(key)
            if lk:
                lk.release()
        except Exception:
            pass

    # ------------------------------------------------------------------
    # Introspection helpers
    # ------------------------------------------------------------------

    @classmethod
    def get_unimplemented_features(cls) -> list[str]:  # NOSONAR - Complexity (19), requires refactoring
        """Return a *copy* of the list enumerating missing capabilities."""
        return list(cls._UNIMPLEMENTED)

    def _extract_protocol_params(self, ca_args: dict) -> tuple:
        """Extract protocol, key_id, and counterparty from create_action args."""
        pd_opts = ca_args.get("pushdrop") or {}
        protocol = (
            ca_args.get("protocol_id")
            or ca_args.get("protocolID")
            or pd_opts.get("protocol_id")
            or pd_opts.get("protocolID")
        )
        key_id = ca_args.get("key_id") or ca_args.get("keyID") or pd_opts.get("key_id") or pd_opts.get("keyID")
        counterparty = ca_args.get("counterparty", pd_opts.get("counterparty"))
        return protocol, key_id, counterparty

    def _normalize_txid(self, txid_val: Any) -> str:
        """Convert txid to hex string format."""
        if isinstance(txid_val, str) and len(txid_val) == 64:
            return txid_val
        elif isinstance(txid_val, (bytes, bytearray)) and len(txid_val) == 32:
            return txid_val.hex()
        else:
            return "00" * 32

    def _create_input_meta(self, output: dict, unlocker: Any, protocol: Any, key_id: Any, counterparty: Any) -> dict:
        """Create metadata for a single input."""
        txid_hex = self._normalize_txid(output.get("txid", ""))
        outpoint = {
            "txid": txid_hex,
            "index": int(output.get("outputIndex", 0)),
        }

        try:
            max_len = unlocker.estimate_length()
        except Exception:
            max_len = 73 + 2

        meta = {
            "outpoint": outpoint,
            "unlockingScriptLength": max_len,
            "inputDescription": output.get("outputDescription", "Previous key-value token"),
            "sequenceNumber": 0,
        }

        # Add optional derived key parameters
        if protocol is not None:
            meta["protocol"] = protocol
        if key_id is not None:
            meta["key_id"] = key_id
        if counterparty is not None:
            meta["counterparty"] = counterparty

        return meta

    def _prepare_inputs_meta(self, key: str, outs: list, ca_args: dict = None) -> list:
        """Prepare the inputs metadata for set/remove operation (Go/TS parity)."""
        ca_args = self._merge_default_ca(ca_args)
        protocol, key_id, counterparty = self._extract_protocol_params(ca_args)

        print(f"[TRACE] [_prepare_inputs_meta] ca_args: {ca_args}")
        print(f"[TRACE] [_prepare_inputs_meta] protocol: {protocol}, key_id: {key_id}, counterparty: {counterparty}")

        pd = PushDrop(self._wallet, self._originator)
        unlock_protocol = protocol if protocol is not None else self._get_protocol(key)
        unlocker = pd.unlock(unlock_protocol, key, None, sign_outputs="all")  # None = SELF counterparty

        inputs_meta = []
        for o in outs:
            meta = self._create_input_meta(o, unlocker, protocol, key_id, counterparty)
            print(f"[TRACE] [_prepare_inputs_meta] meta: {meta}")
            inputs_meta.append(meta)
        return inputs_meta

    def _prepare_spends(self, key, inputs_meta, signable_tx_bytes, input_beef):
        """
        Prepare spends dict for sign_action: {idx: {"unlockingScript": ...}}
        Go/TS parity: use PushDrop unlocker and signable transaction.
        """
        tx = self._link_transaction_for_signing(signable_tx_bytes, input_beef)
        if tx is None:
            return {}

        unlocker = self._create_unlocker(key)
        return self._build_spends_dict(tx, inputs_meta, unlocker)

    def _link_transaction_for_signing(self, signable_tx_bytes, input_beef):
        """Link transaction using BEEF for proper signing context."""
        from bsv.transaction import Transaction, parse_beef_ex
        from bsv.utils import Reader

        try:
            tx = Transaction.from_reader(Reader(signable_tx_bytes))
            if input_beef:
                try:
                    beef, _subject, _last = parse_beef_ex(input_beef)
                    finder = getattr(beef, "find_transaction_for_signing", None)
                    if callable(finder):
                        linked = finder(tx.txid())
                        if linked is not None:
                            tx = linked
                except Exception:
                    pass
            return tx
        except Exception:
            return None

    def _create_unlocker(self, key):
        """Create PushDrop unlocker for the given key."""
        pd = PushDrop(self._wallet, self._originator)
        unlock_protocol = self._get_protocol(key)
        return pd.unlock(unlock_protocol, key, None, sign_outputs="all")  # None = SELF counterparty

    def _build_spends_dict(self, tx, inputs_meta, unlocker):
        """Build spends dictionary for matching inputs."""
        spends = {}
        for idx, meta in enumerate(inputs_meta):
            spend = self._create_spend_for_input(tx, idx, meta, unlocker)
            if spend:
                spends[idx] = spend
        return spends

    def _create_spend_for_input(self, tx, idx, meta, unlocker):
        """Create spend entry for a single input if it matches."""
        try:
            outpoint = meta.get("outpoint") or {}
            meta_txid = outpoint.get("txid")
            meta_index = int(outpoint.get("index", -1))

            # Validate index is within bounds
            if not (0 <= idx < len(tx.inputs)):
                return None

            # Check if outpoint matches
            if not self._outpoint_matches(tx.inputs[idx], meta_txid, meta_index):
                return None

            # Create unlocking script
            unlocking_script = unlocker.sign(tx, idx)
            return {"unlockingScript": unlocking_script}

        except Exception:
            # Skip on error; do not produce empty spends entries
            return None

    def _outpoint_matches(self, tx_input, meta_txid, meta_index):
        """Check if transaction input matches the expected outpoint."""
        try:
            txid_matches = tx_input.source_txid == meta_txid
        except Exception:
            txid_matches = False

        index_matches = getattr(tx_input, "source_output_index", -1) == meta_index
        return txid_matches and index_matches

    # ------------------------------
    # WOC fallback: build minimal BEEF v2
    # ------------------------------
    def _build_beef_v2_from_woc_outputs(self, outputs: list, timeout: int = 10) -> bytes:
        from bsv.beef import build_beef_v2_from_raw_hexes
        from bsv.network.woc_client import WOCClient

        # Collect unique txids present in outputs
        txids: list[str] = []
        for o in outputs:
            txid = o.get("txid")
            if isinstance(txid, str) and len(txid) == 64 and txid != ("00" * 32):
                if txid not in txids:
                    txids.append(txid)
        if not txids:
            return b""
        client = WOCClient()
        tx_hex_list: list[str] = []
        for txid in txids:
            try:
                h = client.get_tx_hex(txid, timeout=timeout)
                if h and isinstance(h, str) and len(h) >= 2:
                    tx_hex_list.append(h)
            except Exception:
                continue
        return build_beef_v2_from_raw_hexes(tx_hex_list)

    def _is_pushdrop_for_pub(self, locking_script_bytes: bytes, pubkey_hex: str | None) -> bool:
        """Rudimentary PushDrop detector: OP_PUSH33 <pubkey33> OP_CHECKSIG then data pushes + DROP.

        This is a heuristic sufficient to filter subject txs for KV get flows.
        """
        try:
            if not pubkey_hex or len(pubkey_hex) != 66:
                return False
            b = locking_script_bytes
            if len(b) < 35:
                return False
            # 0x21 = push 33, followed by 33-byte pubkey, then 0xAC (OP_CHECKSIG)
            if b[0] != 0x21:
                return False
            if b[34] != 0xAC:
                return False
            if b[1:34].hex() != pubkey_hex.lower():
                return False
            # After OP_CHECKSIG must be at least one push and a DROP or 2DROP somewhere
            tail = b[35:]
            if not tail:
                return False
            # Look for OP_DROP(0x75) or OP_2DROP(0x6d)
            return (0x75 in tail) or (0x6D in tail)
        except Exception:
            return False

    # ------------------------------
    # Merge helpers
    # ------------------------------
    def _merge_default_ca(self, ca_args: dict | None) -> dict:
        """Deep-merge config.default_ca into per-call ca_args. ca_args wins.
        Supports nested 'pushdrop' bag similar to TS/GO.
        """
        merged: dict = {}
        # Start with defaults
        if isinstance(self._default_ca, dict):
            merged = copy.deepcopy(self._default_ca)
        # Overlay per-call
        if isinstance(ca_args, dict):
            # top-level scalars
            for k, v in ca_args.items():
                if k == "pushdrop" and isinstance(v, dict):
                    base_pd = merged.get("pushdrop") or {}
                    new_pd = dict(base_pd)
                    new_pd.update(v)
                    merged["pushdrop"] = new_pd
                else:
                    merged[k] = v
        # Ensure feeRate default from config if not set anywhere
        if merged.get("feeRate") is None and merged.get("fee_rate") is None and self._default_fee_rate is not None:
            merged["fee_rate"] = self._default_fee_rate
        return merged


ENTIRE_TXS = "entire transactions"
KEY_EMPTY_MSG = "key cannot be empty"
