from __future__ import annotations

import json
import base64
import hashlib
from dataclasses import dataclass
from typing import Any, Dict, Optional, Tuple, List

import httpx
from ecdsa import SigningKey, SECP256k1
from ecdsa.util import sigencode_string_canonize
from bech32 import bech32_encode, convertbits

from google.protobuf.any_pb2 import Any as ProtoAny
from cosmospy_protobuf.cosmos.tx.v1beta1 import tx_pb2 as cosmos_tx_pb2
from cosmospy_protobuf.cosmos.tx.signing.v1beta1 import signing_pb2 as signing_tx_pb2
from cosmospy_protobuf.cosmos.crypto.secp256k1 import keys_pb2 as secp256k1_keys_pb2
from cosmospy_protobuf.cosmos.base.v1beta1 import coin_pb2 as cosmos_coin_pb2


# ---------------------------------------------------------------------------
# Config / data
# ---------------------------------------------------------------------------

@dataclass
class ChainConfig:
    lcd_url: str
    chain_id: str
    bech32_prefix: str     # "juno", "osmo", "neutron", ...
    base_denom: str        # "ujuno", "uosmo", "untrn", ...


@dataclass
class FeeConfig:
    gas_limit: int = 200_000
    gas_price: int = 1      # in base_denom units per gas (e.g. 1 ujuno / gas)

    @property
    def fee_amount(self) -> int:
        return self.gas_limit * self.gas_price


@dataclass
class Fund:
    denom: str
    amount: str  # string because proto expects string


# ---------------------------------------------------------------------------
# Generic CosmWasm client (execute contract) – ASYNC
# ---------------------------------------------------------------------------

class CosmWasmClient:
    """
    Generic client for CosmWasm MsgExecuteContract using SIGN_MODE_DIRECT.

    Works for any Cosmos SDK chain that:
      - exposes /cosmos/* v1beta1 REST
      - uses secp256k1 accounts
      - has CosmWasm enabled (MsgExecuteContract accepted)
    """

    def __init__(self, cfg: ChainConfig) -> None:
        self.cfg = cfg
        self.lcd_url = cfg.lcd_url.rstrip("/")

    # ----------------- public API -----------------

    async def execute_contract(
        self,
        private_key_hex: str,
        contract_address: str,
        msg: Dict[str, Any],
        funds: Optional[List[Fund]] = None,
        fee_cfg: Optional[FeeConfig] = None,
        memo: str = "",
    ) -> Dict[str, Any]:
        """
        Execute arbitrary CosmWasm contract message.

        - msg: dict -> will be JSON-encoded into MsgExecuteContract.msg
        - funds: optional list of Fund to send coins along

        Returns JSON from /cosmos/tx/v1beta1/txs.
        """
        if fee_cfg is None:
            fee_cfg = FeeConfig()
        if funds is None:
            funds = []

        wallet_addr = self._derive_address_from_private_key(private_key_hex)

        async with httpx.AsyncClient() as client:
            account_number, sequence = await self._get_account_info(client, wallet_addr)

            tx_raw_bytes = self._build_and_sign_tx(
                wallet_address=wallet_addr,
                account_number=account_number,
                sequence=sequence,
                privkey_hex=private_key_hex,
                contract_address=contract_address,
                msg=msg,
                funds=funds,
                fee_cfg=fee_cfg,
                memo=memo,
            )

            return await self._broadcast_tx(client, tx_raw_bytes)

    # ----------------- internal helpers -----------------

    def _derive_address_from_private_key(self, private_key_hex: str) -> str:
        priv_bytes = bytes.fromhex(private_key_hex)
        sk = SigningKey.from_string(priv_bytes, curve=SECP256k1)

        pubkey = sk.get_verifying_key().to_string("compressed")
        sha = hashlib.sha256(pubkey).digest()
        ripe = hashlib.new("ripemd160", sha).digest()

        five_bit = convertbits(ripe, 8, 5, True)
        return bech32_encode(self.cfg.bech32_prefix, five_bit)

    async def _get_account_info(self, client: httpx.AsyncClient, address: str) -> Tuple[int, int]:
        url = f"{self.lcd_url}/cosmos/auth/v1beta1/accounts/{address}"
        r = await client.get(url)
        if r.status_code != 200:
            raise RuntimeError(f"auth query failed {r.status_code}: {r.text}")

        data = r.json()
        if "account" not in data:
            raise RuntimeError(f"no account in response for {address}: {data}")

        account = data["account"]
        base = account.get("base_account", account)
        account_number = int(base["account_number"])
        sequence = int(base["sequence"])
        return account_number, sequence

    # --------- protobuf encoding for MsgExecuteContract ---------

    @staticmethod
    def _encode_varint(value: int) -> bytes:
        out = bytearray()
        v = value
        while True:
            b = v & 0x7F
            v >>= 7
            if v:
                out.append(b | 0x80)
            else:
                out.append(b)
                break
        return bytes(out)

    def _encode_length_delimited(self, field_number: int, payload: bytes) -> bytes:
        key = (field_number << 3) | 2  # wire_type=2 (length-delimited)
        return self._encode_varint(key) + self._encode_varint(len(payload)) + payload

    def _encode_funds(self, funds: List[Fund]) -> bytes:
        """
        Encode repeated cosmos.base.v1beta1.Coin funds = 4;
        message Coin { string denom = 1; string amount = 2; }
        """
        if not funds:
            return b""

        out = bytearray()
        for f in funds:
            coin_bytes = bytearray()
            coin_bytes += self._encode_length_delimited(1, f.denom.encode("utf-8"))
            coin_bytes += self._encode_length_delimited(2, f.amount.encode("utf-8"))

            # field 4 (funds), wire type 2
            out += self._encode_length_delimited(4, bytes(coin_bytes))
        return bytes(out)

    def _build_msg_execute_contract_bytes(
        self,
        sender: str,
        contract_address: str,
        msg: Dict[str, Any],
        funds: List[Fund],
    ) -> bytes:
        """
        Encode cosmwasm.wasm.v1.MsgExecuteContract:

        message MsgExecuteContract {
          string sender   = 1;
          string contract = 2;
          bytes  msg      = 3;
          repeated Coin funds = 4;
        }
        """
        sender_bytes = sender.encode("utf-8")
        contract_bytes = contract_address.encode("utf-8")
        msg_bytes = json.dumps(msg, separators=(",", ":")).encode("utf-8")

        out = bytearray()
        out += self._encode_length_delimited(1, sender_bytes)
        out += self._encode_length_delimited(2, contract_bytes)
        out += self._encode_length_delimited(3, msg_bytes)
        out += self._encode_funds(funds)
        return bytes(out)

    # --------- tx build + sign ---------

    def _build_and_sign_tx(
        self,
        wallet_address: str,
        account_number: int,
        sequence: int,
        privkey_hex: str,
        contract_address: str,
        msg: Dict[str, Any],
        funds: List[Fund],
        fee_cfg: FeeConfig,
        memo: str,
    ) -> bytes:

        # ---- TxBody ----
        msg_bytes = self._build_msg_execute_contract_bytes(
            wallet_address,
            contract_address,
            msg,
            funds,
        )

        msg_any = ProtoAny()
        msg_any.type_url = "/cosmwasm.wasm.v1.MsgExecuteContract"
        msg_any.value = msg_bytes

        tx_body = cosmos_tx_pb2.TxBody(
            messages=[msg_any],
            memo=memo,
            timeout_height=0,
        )
        body_bytes = tx_body.SerializeToString()

        # ---- AuthInfo ----
        priv_bytes = bytes.fromhex(privkey_hex)
        sk = SigningKey.from_string(priv_bytes, curve=SECP256k1)
        pubkey_bytes = sk.get_verifying_key().to_string("compressed")

        pubkey_proto = secp256k1_keys_pb2.PubKey(key=pubkey_bytes)
        pubkey_any = ProtoAny()
        pubkey_any.Pack(pubkey_proto)
        pubkey_any.type_url = "/cosmos.crypto.secp256k1.PubKey"

        mode_info = cosmos_tx_pb2.ModeInfo(
            single=cosmos_tx_pb2.ModeInfo.Single(
                mode=signing_tx_pb2.SIGN_MODE_DIRECT
            )
        )

        signer_info = cosmos_tx_pb2.SignerInfo(
            public_key=pubkey_any,
            mode_info=mode_info,
            sequence=sequence,
        )

        fee = cosmos_tx_pb2.Fee(
            amount=[
                cosmos_coin_pb2.Coin(
                    denom=self.cfg.base_denom,
                    amount=str(fee_cfg.fee_amount),
                )
            ],
            gas_limit=fee_cfg.gas_limit,
            payer="",
            granter="",
        )

        auth_info = cosmos_tx_pb2.AuthInfo(
            signer_infos=[signer_info],
            fee=fee,
        )
        auth_info_bytes = auth_info.SerializeToString()

        # ---- SignDoc ----
        sign_doc = cosmos_tx_pb2.SignDoc(
            body_bytes=body_bytes,
            auth_info_bytes=auth_info_bytes,
            chain_id=self.cfg.chain_id,
            account_number=account_number,
        )

        sign_doc_bytes = sign_doc.SerializeToString()
        sign_doc_hash = hashlib.sha256(sign_doc_bytes).digest()

        signature_bytes = sk.sign_digest(
            sign_doc_hash,
            sigencode=sigencode_string_canonize,
        )

        # ---- TxRaw ----
        tx_raw = cosmos_tx_pb2.TxRaw(
            body_bytes=body_bytes,
            auth_info_bytes=auth_info_bytes,
            signatures=[signature_bytes],
        )

        return tx_raw.SerializeToString()

    # --------- broadcast ---------

    async def _broadcast_tx(self, client: httpx.AsyncClient, tx_raw_bytes: bytes) -> Dict[str, Any]:
        tx_bytes_b64 = base64.b64encode(tx_raw_bytes).decode("utf-8")
        payload = {
            "tx_bytes": tx_bytes_b64,
            "mode": "BROADCAST_MODE_SYNC",
        }
        url = f"{self.lcd_url}/cosmos/tx/v1beta1/txs"
        r = await client.post(url, json=payload)
        if r.status_code != 200:
            raise RuntimeError(f"broadcast failed {r.status_code}: {r.text}")
        return r.json()
