from dataclasses import dataclass
import logging
import random
import time
from typing import Any, Dict, List, Optional, TypeVar
from eth_account import Account
from eth_account.datastructures import SignedMessage, SignedTransaction
from eth_account.signers.local import LocalAccount
from eth_account.messages import encode_typed_data
from pydantic import BaseModel, ConfigDict, Field
from pydantic.alias_generators import to_camel
from web3 import AsyncBaseProvider, AsyncWeb3
from web3.contract import AsyncContract
from functools import cache, lru_cache
import json
from importlib import resources as impresources

from web3 import Web3
from web3.eth import AsyncEth
from eth_utils import to_checksum_address

from tristero.api import get_quote

from .data import (
    get_permit2_addr,
    get_wrapped_gas_addr,
)
from .eip712 import (
    EIP712OrderParameters,
    EIP712TokenPermissions,
    EIP712SignedOrder,
    EIP712PermitWitnessTransferFromOrder,
    EIP712CloseWithSwap,
    get_escrow_domain,
    get_close_with_swap_types,
)

logger = logging.getLogger(__name__)

P = TypeVar("P", bound=AsyncBaseProvider)

DEFAULT_PERMIT2_ADDRESS = "0x000000000022D473030F116dDEE9F6B43aC78BA3"

PERMIT2_ABI_FILE = impresources.files("tristero.files") / "permit2_abi.json"
ERC20_ABI_FILE = impresources.files("tristero.files") / "erc20_abi.json"

PERMIT2_ABI = json.loads(PERMIT2_ABI_FILE.read_text())
ERC20_ABI = json.loads(ERC20_ABI_FILE.read_text())


@lru_cache(maxsize=None)
def get_permit2_contract(eth: AsyncEth, permit2_address: str):
    return eth.contract(
        address=Web3.to_checksum_address(permit2_address), abi=PERMIT2_ABI
    )


@cache
def get_erc20_contract(w3: AsyncWeb3[P], token_address: str) -> AsyncContract:
    """Get ERC20 contract instance."""
    return w3.eth.contract(
        address=Web3.to_checksum_address(token_address), abi=ERC20_ABI
    )


class BaseSchema(BaseModel, frozen=True):
    model_config = ConfigDict(alias_generator=to_camel, populate_by_name=True)


class TokenData(BaseSchema, frozen=True):
    address: str
    value: int


class ChainData(BaseSchema, frozen=True):
    chain_id: int
    token: TokenData


class OrderParameters(BaseSchema, frozen=True):
    src_asset: str
    dst_asset: str
    src_quantity: str
    dst_quantity: str
    min_quantity: str
    dark_salt: str


class OrderData(BaseSchema, frozen=True):
    parameters: OrderParameters
    deadline: int
    router_address: str
    filler_wallet_address: str
    order_type: str
    custom_data: List[bytes] = Field(default_factory=list)


class Quote(BaseSchema, frozen=True):
    order_data: OrderData


class SignedOrder(BaseSchema, frozen=True):
    sender: str
    parameters: OrderParameters
    deadline: str
    target: str
    filler: str
    order_type: str
    custom_data: List[bytes]


class TokenPermissions(BaseSchema, frozen=True):
    token: str
    amount: str


class EIP712Domain(BaseSchema, frozen=True):
    name: str
    chain_id: int
    verifying_contract: str


class PermitMessage(BaseSchema, frozen=True):
    permitted: TokenPermissions
    spender: str
    nonce: str
    deadline: str
    witness: SignedOrder


class SignatureData(BaseSchema, frozen=True):
    domain: EIP712Domain
    types: dict[str, Any]
    primary_type: str
    message: PermitMessage


def first_zero_bit(n: int) -> Optional[int]:
    if n > 0 and (n & (n + 1)) == 0:
        return None
    return (n ^ (n + 1)).bit_length() - 1


async def get_permit2_unordered_nonce(c: AsyncContract, wallet_address: str):
    startWord = random.randint(2**126, 2**189)
    endWord = startWord + 10
    wordPos = await anext(
        (
            ((w << 8) | bit)
            for w in range(startWord, endWord)
            if (
                bit := first_zero_bit(
                    await c.functions.nonceBitmap(wallet_address, w).call()
                )
            )
            is not None
        ),
        None,
    )
    if wordPos is None:
        raise Exception(f"No free unordered nonces in words {startWord}-{endWord}")
    return wordPos


def get_random_nonce() -> int:
    """Generate a random nonce for EIP-712 signing."""
    return random.randint(2**128, 2**256 - 1)


async def prepare_data_for_signature(
    eth: AsyncEth,
    sell_data: ChainData,
    buy_data: ChainData,
    wallet_address: str,
    quote: Quote,
    destination_address: Optional[str] = None,
) -> SignatureData:
    """
    Prepare EIP-712 signature data for Permit2 witness transfer.

    Args:
        sell_data: Source chain and token data
        buy_data: Destination chain and token data
        wallet_address: User's wallet address
        quote: Quote data with order parameters and deadline
        destination_address: Optional destination address

    Returns:
        SignatureData with domain, types, primaryType, and message for signing
    """
    if not quote.order_data.parameters.min_quantity:
        raise ValueError("Min quantity is required in the order_data.parameters")

    if not quote.order_data.router_address:
        raise ValueError("Router address is required in the order_data")

    if not quote.order_data.deadline:
        raise ValueError("Deadline is required in the order_data")

    from_chain = str(sell_data.chain_id)

    deadline = quote.order_data.deadline

    token_address = (
        get_wrapped_gas_addr(from_chain)
        if sell_data.token.address == "native"
        else sell_data.token.address
    )

    witness = SignedOrder(
        sender=wallet_address,
        parameters=OrderParameters(
            src_asset=token_address,
            dst_asset=buy_data.token.address,
            src_quantity=str(sell_data.token.value),
            dst_quantity=str(buy_data.token.value),
            min_quantity=quote.order_data.parameters.min_quantity,
            dark_salt=quote.order_data.parameters.dark_salt,
        ),
        deadline=str(deadline),
        target=destination_address or wallet_address,
        filler=quote.order_data.filler_wallet_address,
        order_type=quote.order_data.order_type,
        custom_data=quote.order_data.custom_data or [],
    )

    permit2_address = get_permit2_addr(from_chain)
    if not permit2_address:
        raise ValueError("Permit2 not deployed on this chain.")

    spender = quote.order_data.router_address
    nonce = await get_permit2_unordered_nonce(
        get_permit2_contract(eth, permit2_address), wallet_address
    )

    domain = EIP712Domain(
        name="Permit2",
        chain_id=sell_data.chain_id,
        verifying_contract=permit2_address,
    )

    types = {
        "TokenPermissions": [
            {"name": "token", "type": "address"},
            {"name": "amount", "type": "uint256"},
        ],
        "OrderParameters": [
            {"name": "srcAsset", "type": "address"},
            {"name": "dstAsset", "type": "address"},
            {"name": "srcQuantity", "type": "uint256"},
            {"name": "dstQuantity", "type": "uint256"},
            {"name": "minQuantity", "type": "uint256"},
            {"name": "darkSalt", "type": "uint128"},
        ],
        "SignedOrder": [
            {"name": "sender", "type": "address"},
            {"name": "parameters", "type": "OrderParameters"},
            {"name": "deadline", "type": "uint256"},
            {"name": "target", "type": "address"},
            {"name": "filler", "type": "address"},
            {"name": "orderType", "type": "string"},
            {"name": "customData", "type": "bytes[]"},
        ],
        "PermitWitnessTransferFrom": [
            {"name": "permitted", "type": "TokenPermissions"},
            {"name": "spender", "type": "address"},
            {"name": "nonce", "type": "uint256"},
            {"name": "deadline", "type": "uint256"},
            {"name": "witness", "type": "SignedOrder"},
        ],
    }

    message = PermitMessage(
        permitted=TokenPermissions(
            token=token_address,
            amount=str(sell_data.token.value),
        ),
        spender=spender,
        nonce=str(nonce),
        deadline=str(deadline),
        witness=witness,
    )

    return SignatureData(
        domain=domain,
        types=types,
        primary_type="PermitWitnessTransferFrom",
        message=message,
    )


async def sign_permit2(
    eth: AsyncEth,
    account: LocalAccount,
    wallet_address: str,
    src_amount: int,
    quote_data: dict[str, Any],
) -> tuple[SignatureData, SignedMessage]:
    from_chain_id = quote_data["src_token"]["chain_id"]
    from_address = quote_data["src_token"]["address"]
    to_chain_id = quote_data["dst_token"]["chain_id"]
    to_address = quote_data["dst_token"]["address"]

    order_data = quote_data["order_data"]
    dst_amount = int(order_data["parameters"]["dst_quantity"])

    to_sign = await prepare_data_for_signature(
        eth,
        ChainData(
            chain_id=from_chain_id,
            token=TokenData(address=from_address, value=src_amount),
        ),
        ChainData(
            chain_id=to_chain_id,
            token=TokenData(address=to_address, value=dst_amount),
        ),
        wallet_address,
        Quote(order_data=order_data),
    )
    signature = account.sign_typed_data(
        full_message=to_sign.model_dump(mode="json", by_alias=True)
    )
    return (to_sign, signature)


def sign_margin_order(
    quote: Dict[str, Any],
    private_key: str,
    permit2_address: str = DEFAULT_PERMIT2_ADDRESS,
    nonce: Optional[int] = None,
) -> Dict[str, Any]:
    """
    Sign a margin order using EIP-712 typed data.
    
    Args:
        quote: Quote response from get_margin_quote
        private_key: Private key for signing
        permit2_address: Permit2 contract address
        nonce: Optional nonce (random if not provided)
    
    Returns:
        Signed order payload ready for submission
    """
    order_data = quote.get("order_data")
    if not isinstance(order_data, dict):
        raise ValueError("quote missing order_data")

    chain_id_str = quote.get("chain_id")
    if isinstance(chain_id_str, dict) and "value" in chain_id_str:
        chain_id = int(chain_id_str["value"])
    else:
        chain_id = int(chain_id_str)

    domain_data = {
        "name": "Permit2",
        "chainId": chain_id,
        "verifyingContract": to_checksum_address(permit2_address),
    }

    parameters = EIP712OrderParameters(
        srcAsset=order_data["parameters"]["src_asset"],
        dstAsset=order_data["parameters"]["dst_asset"],
        srcQuantity=int(order_data["parameters"]["src_quantity"]),
        dstQuantity=int(order_data["parameters"]["dst_quantity"]),
        minQuantity=int(order_data["parameters"]["min_quantity"]),
        darkSalt=int(order_data["parameters"]["dark_salt"]),
    )

    signed_order = EIP712SignedOrder(
        sender=order_data["sender_wallet_address"],
        parameters=parameters,
        deadline=int(order_data["deadline"]),
        target=order_data["target_wallet_address"],
        filler=order_data["filler_wallet_address"],
        orderType=str(order_data["order_type"]),
        customData=[
            bytes.fromhex(cd[2:]) if cd.startswith("0x") else bytes.fromhex(cd)
            for cd in order_data.get("custom_data", [])
        ],
    )

    token_permissions = EIP712TokenPermissions(
        token=order_data["parameters"]["src_asset"],
        amount=int(order_data["parameters"]["src_quantity"]),
    )

    if nonce is None:
        nonce = get_random_nonce()

    permit_witness = EIP712PermitWitnessTransferFromOrder(
        permitted=token_permissions,
        spender=order_data["router_address"],
        nonce=nonce,
        deadline=int(order_data["deadline"]),
        witness=signed_order,
    )

    typed_data = {
        "types": {
            "EIP712Domain": [
                {"name": "name", "type": "string"},
                {"name": "chainId", "type": "uint256"},
                {"name": "verifyingContract", "type": "address"},
            ],
            "PermitWitnessTransferFrom": permit_witness.TYPE_STRUCT,
            "TokenPermissions": token_permissions.TYPE_STRUCT,
            "SignedOrder": signed_order.TYPE_STRUCT,
            "OrderParameters": parameters.TYPE_STRUCT,
        },
        "primaryType": "PermitWitnessTransferFrom",
        "domain": domain_data,
        "message": permit_witness.eip_signable_struct,
    }

    acct = Account.from_key(private_key)
    signable = encode_typed_data(full_message=typed_data)
    sig = acct.sign_message(signable).signature.hex()
    if not sig.startswith("0x"):
        sig = "0x" + sig

    return {
        "signature": sig,
        "domain": domain_data,
        "message": permit_witness.eip_signable_struct,
    }


def sign_close_position(
    chain_id: int,
    position_id: int,
    private_key: str,
    escrow_contract: str,
    authorized: str,
    cash_settle: bool = False,
    fraction_bps: int = 10_000,
    deadline_seconds: int = 3600,
    nonce: Optional[int] = None,
) -> Dict[str, Any]:
    """
    Sign a close margin position request using EIP-712.
    
    Args:
        chain_id: Chain ID
        position_id: Position ID (NFT token ID)
        private_key: Private key for signing
        escrow_contract: Escrow contract address
        authorized: Authorized filler address
        cash_settle: Whether to cash settle
        fraction_bps: Fraction to close in basis points (10000 = 100%)
        deadline_seconds: Deadline in seconds from now
        nonce: Optional nonce (random if not provided)
    
    Returns:
        Signed close position payload
    """
    if not escrow_contract:
        raise ValueError("escrow_contract is required")
    if not authorized:
        raise ValueError("authorized is required")

    if nonce is None:
        nonce = get_random_nonce()

    deadline = int(time.time()) + int(deadline_seconds)

    close_msg = EIP712CloseWithSwap(
        positionId=int(position_id),
        cashSettle=bool(cash_settle),
        fractionBps=int(fraction_bps),
        authorized=str(authorized),
        nonce=int(nonce),
        deadline=int(deadline),
    )

    domain_data = get_escrow_domain(chain_id, escrow_contract)
    types = get_close_with_swap_types()

    typed_data = {
        "types": types,
        "primaryType": "CloseWithSwap",
        "domain": domain_data,
        "message": {
            "positionId": close_msg.positionId,
            "cashSettle": close_msg.cashSettle,
            "fractionBps": close_msg.fractionBps,
            "authorized": close_msg.authorized,
            "nonce": close_msg.nonce,
            "deadline": close_msg.deadline,
        },
    }

    acct = Account.from_key(private_key)
    signable = encode_typed_data(full_message=typed_data)
    sig = acct.sign_message(signable).signature.hex()
    if not sig.startswith("0x"):
        sig = "0x" + sig

    return {
        "signature": sig,
        "domain": domain_data,
        "message": typed_data["message"],
        "chainId": str(chain_id),
    }


async def approve_permit2(
    w3: AsyncWeb3[P],
    account: LocalAccount,
    chain: str,
    token_address: str,
    required_quantity: int,
    maxGas: int = 100000,
):
    wallet_address = account.address
    permit2_address = get_permit2_addr(chain)

    erc20 = get_erc20_contract(w3, token_address)
    current_allowance = await erc20.functions.allowance(
        wallet_address, permit2_address
    ).call()
    if current_allowance < required_quantity:
        logger.info(
            f"Approving {token_address}: allowance={current_allowance}, required={required_quantity}"
        )
        approve_fn = erc20.functions.approve(permit2_address, 2**256 - 1)
        tx = await approve_fn.build_transaction(
            {
                "from": wallet_address,
                "nonce": await w3.eth.get_transaction_count(
                    w3.to_checksum_address(wallet_address)
                ),
                "gas": maxGas,
                "gasPrice": await w3.eth.gas_price,
            }
        )

        tx_data = tx if isinstance(tx, dict) else tx.__dict__
        signed_tx: SignedTransaction = account.sign_transaction(tx_data)
        tx_hash = await w3.eth.send_raw_transaction(signed_tx.raw_transaction)

        logger.debug(f"→ Approval tx hash: {tx_hash.hex()}")
        return tx_hash.hex()


@dataclass
class Permit2Order:
    msg: SignatureData
    sig: SignedMessage


async def create_permit2_order(
    w3: AsyncWeb3[P],
    account: LocalAccount,
    src_chain: str,
    src_token: str,
    dst_chain: str,
    dst_token: str,
    raw_amount: int,
    dst_address: str | None = None,
) -> Permit2Order:
    if not dst_address:
        dst_address = account.address
    q = await get_quote(
        account.address,
        dst_address,
        src_chain,
        src_token,
        dst_chain,
        dst_token,
        raw_amount,
    )

    order_type = q.get("orderType") or q.get("order_type")
    if not order_type:
        order_data = q.get("order_data")
        if isinstance(order_data, dict):
            order_type = order_data.get("order_type") or order_data.get("orderType")

    if not order_type:
        raise ValueError(f"quote missing order type (keys={sorted(q.keys())})")

    if str(order_type).strip().upper() == "FEATHER":
        raise Exception("Feather routes unsupported for automatic trades")

    _ = await approve_permit2(w3, account, src_chain, src_token, raw_amount)
    msg, sig = await sign_permit2(w3.eth, account, account.address, raw_amount, q)
    return Permit2Order(msg, sig)
