"""
StarkEx signing adapter for the EdgeX Python SDK.

This module provides an implementation of the signing adapter interface
that uses the StarkWare cryptographic primitives for signing operations.
"""

import binascii
import math
import secrets
from typing import List, Tuple

from .signing_adapter import SigningAdapter
from ..crypto.pedersen_hash import pedersen_hash_bytes


# StarkEx curve parameters
FIELD_PRIME = 0x800000000000011000000000000000000000000000000000000000000000001
ALPHA = 1
BETA = 0x6f21413efbe40de150e596d72f7a8c5609ad26c15c915c1f4cdfcb99cee9e89
EC_ORDER = 0x800000000000010ffffffffffffffffb781126dcae7b2321e66a241adc64d2f
N_ELEMENT_BITS_ECDSA = math.floor(math.log(FIELD_PRIME, 2))
assert N_ELEMENT_BITS_ECDSA == 251

# Generator point for the Stark curve
EC_GEN = (
    0x1ef15c18599971b7beced415a40f0c7deacfd9b0d1819e03d723d8bc943cfca,
    0x5668060aa49730b7be4801df46ec62de53ecd11abe43a32873000c36e8dc1f
)


class StarkExSigningAdapter(SigningAdapter):
    """StarkEx implementation of the signing adapter interface."""

    def sign(self, message_hash: bytes, private_key: str) -> Tuple[str, str]:
        """
        Sign a message hash using a private key.

        Args:
            message_hash: The hash of the message to sign
            private_key: The private key as a hex string

        Returns:
            Tuple[str, str]: The signature as (r, s) hex strings

        Raises:
            ValueError: If the private key is invalid or the signing fails
        """
        try:
            # Validate private key format
            binascii.unhexlify(private_key)
        except binascii.Error:
            raise ValueError("Invalid private key hex string")

        # Convert message hash to integer
        msg_hash_int = int.from_bytes(message_hash, byteorder='big')

        # Ensure the message hash is in the valid range
        # Use the same modulus as the Golang SDK (EC_ORDER, which is starkcurve.N)
        msg_hash_int = msg_hash_int % EC_ORDER

        # Convert private key to integer
        priv_key_int = int(private_key, 16)

        # Ensure the private key is in the valid range
        # For testing purposes, we'll just take the modulus
        priv_key_int = priv_key_int % EC_ORDER
        if priv_key_int == 0:
            priv_key_int = 1

        # Sign the message
        r, s = self._sign(msg_hash_int, priv_key_int)

        # Convert r and s to hex strings
        r_hex = format(r, '064x')
        s_hex = format(s, '064x')

        return r_hex, s_hex

    def get_public_key(self, private_key: str) -> str:
        """
        Get the public key from a private key.

        Args:
            private_key: The private key as a hex string

        Returns:
            str: The public key as a hex string

        Raises:
            ValueError: If the private key is invalid
        """
        try:
            # Validate private key format
            binascii.unhexlify(private_key)
        except binascii.Error:
            raise ValueError("Invalid private key hex string")

        # Convert private key to integer
        priv_key_int = int(private_key, 16)

        # Ensure the private key is in the valid range
        # For testing purposes, we'll just take the modulus
        priv_key_int = priv_key_int % EC_ORDER
        if priv_key_int == 0:
            priv_key_int = 1

        # Get the public key
        public_key = self._private_to_stark_key(priv_key_int)

        # Convert public key to hex string
        public_key_hex = format(public_key, '064x')

        return public_key_hex

    def verify(self, message_hash: bytes, signature: Tuple[str, str], public_key: str) -> bool:
        """
        Verify a signature using a public key.

        Args:
            message_hash: The hash of the message
            signature: The signature as (r, s) hex strings
            public_key: The public key as a hex string

        Returns:
            bool: Whether the signature is valid
        """
        try:
            # Convert message hash to integer
            msg_hash_int = int.from_bytes(message_hash, byteorder='big')

            # Ensure the message hash is in the valid range
            # Use the same modulus as the sign method (EC_ORDER)
            msg_hash_int = msg_hash_int % EC_ORDER

            # Convert signature components to integers
            r_int = int(signature[0], 16)
            s_int = int(signature[1], 16)

            # Ensure r and s are in the valid range
            if not (1 <= r_int < 2**N_ELEMENT_BITS_ECDSA and 1 <= s_int < EC_ORDER):
                return False

            # Convert public key to integer
            pub_key_int = int(public_key, 16)

            # Verify the signature
            return self._verify(msg_hash_int, r_int, s_int, pub_key_int)
        except Exception:
            return False

    def pedersen_hash(self, elements: List[int]) -> bytes:
        """
        Calculate the Pedersen hash of a list of integers.

        This method now uses the full Pedersen hash implementation
        that follows StarkWare's specification.

        Args:
            elements: List of integers to hash

        Returns:
            bytes: The hash result

        Raises:
            ValueError: If the calculation fails
        """
        try:
            # Use the full Pedersen hash implementation
            return pedersen_hash_bytes(*elements)
        except Exception as e:
            raise ValueError(f"Failed to calculate Pedersen hash: {str(e)}")

    def _sign(self, msg_hash: int, priv_key: int) -> Tuple[int, int]:
        """
        Sign a message hash using a private key.

        Args:
            msg_hash: The hash of the message to sign as an integer
            priv_key: The private key as an integer

        Returns:
            Tuple[int, int]: The signature as (r, s) integers
        """
        # Choose a valid k. In our version of ECDSA not every k value is valid,
        # and there is a negligible probability a drawn k cannot be used for signing.
        # This is why we have this loop.
        while True:
            # Use random nonce generation like the Go SDK
            k = self._generate_random_k()

            # Cannot fail because 0 < k < EC_ORDER and EC_ORDER is prime.
            x = self._ec_mult(k, EC_GEN)[0]

            # DIFF: in classic ECDSA, we take int(x) % n.
            r = int(x)
            if not (1 <= r < 2**N_ELEMENT_BITS_ECDSA):
                # Bad value. This fails with negligible probability.
                continue

            if (msg_hash + r * priv_key) % EC_ORDER == 0:
                # Bad value. This fails with negligible probability.
                continue

            w = self._div_mod(k, msg_hash + r * priv_key, EC_ORDER)
            if not (1 <= w < 2**N_ELEMENT_BITS_ECDSA):
                # Bad value. This fails with negligible probability.
                continue

            s = self._inv_mod_curve_size(w)
            return r, s

    def _verify(self, msg_hash: int, r: int, s: int, public_key: int) -> bool:
        """
        Verify a signature using a public key.

        Args:
            msg_hash: The hash of the message as an integer
            r: The r component of the signature as an integer
            s: The s component of the signature as an integer
            public_key: The public key as an integer

        Returns:
            bool: Whether the signature is valid
        """
        # Compute w = s^-1 (mod EC_ORDER).
        if not (1 <= s < EC_ORDER):
            return False

        w = self._inv_mod_curve_size(s)

        # Preassumptions:
        # DIFF: in classic ECDSA, we assert 1 <= r, w <= EC_ORDER-1.
        # Since r, w < 2**N_ELEMENT_BITS_ECDSA < EC_ORDER, we only need to verify r, w != 0.
        if not (1 <= r < 2**N_ELEMENT_BITS_ECDSA and 1 <= w < 2**N_ELEMENT_BITS_ECDSA):
            return False

        if not (0 <= msg_hash < 2**N_ELEMENT_BITS_ECDSA):
            return False

        # Only the x coordinate of the point is given, check the two possibilities for the y
        # coordinate.
        try:
            y = self._get_y_coordinate(public_key)
        except ValueError:
            return False

        # Verify it is on the curve.
        if (y**2 - (public_key**3 + ALPHA * public_key + BETA)) % FIELD_PRIME != 0:
            return False

        # Try both possible y coordinates.
        for y_candidate in [y, (-y) % FIELD_PRIME]:
            public_key_point = (public_key, y_candidate)

            # Signature validation.
            try:
                # Calculate u1 = msg_hash * w mod n
                u1 = (msg_hash * w) % EC_ORDER

                # Calculate u2 = r * w mod n
                u2 = (r * w) % EC_ORDER

                # Calculate u1*G + u2*Q
                point1 = self._ec_mult(u1, EC_GEN)
                point2 = self._ec_mult(u2, public_key_point)
                point = self._ec_add(point1, point2)

                # The signature is valid if the x-coordinate of the resulting point equals r
                if point[0] == r:
                    return True
            except Exception:
                continue

        return False

    def _generate_random_k(self) -> int:
        """
        Generate a cryptographically secure random k value.

        Returns:
            int: The generated k value in range [1, EC_ORDER)
        """
        # Generate a cryptographically secure random number in the range [1, EC_ORDER)
        # This matches the Go implementation's approach of using random nonces
        return secrets.randbelow(EC_ORDER - 1) + 1

    def _private_to_stark_key(self, priv_key: int) -> int:
        """
        Convert a private key to a Stark public key.

        Args:
            priv_key: The private key as an integer

        Returns:
            int: The public key as an integer
        """
        return self._private_key_to_ec_point_on_stark_curve(priv_key)[0]

    def _private_key_to_ec_point_on_stark_curve(self, priv_key: int) -> Tuple[int, int]:
        """
        Convert a private key to an EC point on the Stark curve.

        Args:
            priv_key: The private key as an integer

        Returns:
            Tuple[int, int]: The EC point as (x, y) coordinates
        """
        # Ensure the private key is in the valid range
        # For testing purposes, we'll just take the modulus
        priv_key = priv_key % EC_ORDER
        if priv_key == 0:
            priv_key = 1

        return self._ec_mult(priv_key, EC_GEN)

    def _inv_mod_curve_size(self, x: int) -> int:
        """
        Calculate the modular inverse of x modulo the curve order.

        Args:
            x: The value to invert

        Returns:
            int: The modular inverse
        """
        return self._div_mod(1, x, EC_ORDER)

    def _div_mod(self, n: int, m: int, p: int) -> int:
        """
        Calculate (n / m) mod p.

        Args:
            n: The numerator
            m: The denominator
            p: The modulus

        Returns:
            int: The result of the division modulo p
        """
        return (n * pow(m, -1, p)) % p

    def _is_quad_residue(self, n: int, p: int) -> bool:
        """
        Check if n is a quadratic residue modulo p.

        Args:
            n: The number to check
            p: The modulus

        Returns:
            bool: True if n is a quadratic residue modulo p, False otherwise
        """
        return pow(n, (p - 1) // 2, p) == 1

    def _sqrt_mod(self, n: int, p: int) -> int:
        """
        Calculate the square root of n modulo p.

        Args:
            n: The number to take the square root of
            p: The modulus

        Returns:
            int: The square root of n modulo p
        """
        # Handle the case where p = 3 mod 4
        if p % 4 == 3:
            return pow(n, (p + 1) // 4, p)

        # Handle the general case using the Tonelli-Shanks algorithm
        q = p - 1
        s = 0
        while q % 2 == 0:
            q //= 2
            s += 1

        # Find a non-residue
        z = 2
        while self._is_quad_residue(z, p):
            z += 1

        m = s
        c = pow(z, q, p)
        t = pow(n, q, p)
        r = pow(n, (q + 1) // 2, p)

        while t != 1:
            # Find the least i, 0 < i < m, such that t^(2^i) = 1
            i = 0
            t_sq = t
            while t_sq != 1 and i < m - 1:
                t_sq = (t_sq * t_sq) % p
                i += 1

            # Calculate b = c^(2^(m-i-1))
            b = pow(c, 2**(m - i - 1), p)

            m = i
            c = (b * b) % p
            t = (t * b * b) % p
            r = (r * b) % p

        return r

    def _get_y_coordinate(self, x: int) -> int:
        """
        Given the x coordinate of a point, returns a possible y coordinate such that
        together the point (x,y) is on the curve.

        Args:
            x: The x coordinate

        Returns:
            int: A possible y coordinate

        Raises:
            ValueError: If x is not a valid x coordinate on the curve
        """
        y_squared = (x * x * x + ALPHA * x + BETA) % FIELD_PRIME
        if not self._is_quad_residue(y_squared, FIELD_PRIME):
            raise ValueError("Given x coordinate does not represent any point on the elliptic curve.")

        return self._sqrt_mod(y_squared, FIELD_PRIME)

    def _ec_add(self, p1: Tuple[int, int], p2: Tuple[int, int]) -> Tuple[int, int]:
        """
        Add two points on the elliptic curve.

        Args:
            p1: The first point as (x, y) coordinates
            p2: The second point as (x, y) coordinates

        Returns:
            Tuple[int, int]: The resulting point as (x, y) coordinates
        """
        if p1[0] == p2[0]:
            if (p1[1] + p2[1]) % FIELD_PRIME == 0:
                # The points are negatives of each other, return the point at infinity
                # We represent the point at infinity as None, but this should never happen
                # in our use case, so we raise an exception instead
                raise ValueError("Points are negatives of each other")

            # The points are the same, so we're doubling
            return self._ec_double(p1)

        # Calculate the slope
        slope = self._div_mod(p2[1] - p1[1], p2[0] - p1[0], FIELD_PRIME)

        # Calculate the new point
        x3 = (slope * slope - p1[0] - p2[0]) % FIELD_PRIME
        y3 = (slope * (p1[0] - x3) - p1[1]) % FIELD_PRIME

        return (x3, y3)

    def _ec_double(self, p: Tuple[int, int]) -> Tuple[int, int]:
        """
        Double a point on the elliptic curve.

        Args:
            p: The point to double as (x, y) coordinates

        Returns:
            Tuple[int, int]: The resulting point as (x, y) coordinates
        """
        # Calculate the slope
        slope = self._div_mod(3 * p[0] * p[0] + ALPHA, 2 * p[1], FIELD_PRIME)

        # Calculate the new point
        x3 = (slope * slope - 2 * p[0]) % FIELD_PRIME
        y3 = (slope * (p[0] - x3) - p[1]) % FIELD_PRIME

        return (x3, y3)

    def _ec_mult(self, m: int, p: Tuple[int, int]) -> Tuple[int, int]:
        """
        Multiply a point on the elliptic curve by a scalar.

        Args:
            m: The scalar
            p: The point as (x, y) coordinates

        Returns:
            Tuple[int, int]: The resulting point as (x, y) coordinates
        """
        if m == 0:
            raise ValueError("Cannot multiply by 0")

        if m == 1:
            return p

        if m % 2 == 0:
            return self._ec_mult(m // 2, self._ec_double(p))
        else:
            return self._ec_add(p, self._ec_mult(m - 1, p))