import os
from typing import Tuple
from cryptography.exceptions import InvalidSignature
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import ec, utils
from cryptography.hazmat.primitives.asymmetric.utils import decode_dss_signature, encode_dss_signature

def generate_identity_keys() -> Tuple[ec.EllipticCurvePrivateKey, bytes]:
    """
    Generate an ECDSA (P-256) identity key pair.
    
    Returns:
        Tuple containing:
        - private_key: EllipticCurvePrivateKey object
        - public_key_bytes: Raw x||y format (64 bytes, 512 bits)
    """
    private_key = ec.generate_private_key(ec.SECP256R1())
    public_numbers = private_key.public_key().public_numbers()

    x_bytes = public_numbers.x.to_bytes(32, byteorder='big')
    y_bytes = public_numbers.y.to_bytes(32, byteorder='big')
    pubkey_bytes = x_bytes + y_bytes  # 64 bytes total

    return private_key, pubkey_bytes


def load_peer_identity_key(pubkey_bytes: bytes) -> ec.EllipticCurvePublicKey:
    """
    Convert a raw public key (64 bytes, x||y format) to a cryptography public key object.
    
    Args:
        pubkey_bytes: Raw 64-byte public key (x||y format)
        
    Returns:
        EllipticCurvePublicKey object
        
    Raises:
        ValueError: If the pubkey_bytes is not exactly 64 bytes
    """
    if len(pubkey_bytes) != 64:
        raise ValueError("Peer identity pubkey must be exactly 64 bytes (x||y).")
        
    x_int = int.from_bytes(pubkey_bytes[:32], byteorder='big')
    y_int = int.from_bytes(pubkey_bytes[32:], byteorder='big')
    
    public_numbers = ec.EllipticCurvePublicNumbers(x_int, y_int, ec.SECP256R1())
    return public_numbers.public_key()


def sign_data(private_key: ec.EllipticCurvePrivateKey, data: bytes) -> bytes:
    """
    Sign data with ECDSA using a P-256 private key.
    
    Args:
        private_key: EllipticCurvePrivateKey for signing
        data: Bytes to sign
        
    Returns:
        DER-encoded signature (variable length, up to ~70-72 bytes)
    """
    signature = private_key.sign(data, ec.ECDSA(hashes.SHA256()))
    return signature


def verify_signature(public_key: ec.EllipticCurvePublicKey, signature: bytes, data: bytes) -> bool:
    """
    Verify a DER-encoded ECDSA signature.
    
    Args:
        public_key: EllipticCurvePublicKey for verification
        signature: DER-encoded signature
        data: Original signed data
        
    Returns:
        True if signature is valid, False otherwise
    """
    try:
        public_key.verify(signature, data, ec.ECDSA(hashes.SHA256()))
        return True
    except InvalidSignature:
        return False


def get_ephemeral_keypair() -> Tuple[ec.EllipticCurvePrivateKey, bytes]:
    """
    Generate an ephemeral ECDH key pair (P-256).
    
    Returns:
        Tuple containing:
        - private_key: EllipticCurvePrivateKey object
        - pubkey_bytes: Raw x||y format (64 bytes, 512 bits)
    """
    private_key = ec.generate_private_key(ec.SECP256R1())
    numbers = private_key.public_key().public_numbers()
    
    x_bytes = numbers.x.to_bytes(32, 'big')
    y_bytes = numbers.y.to_bytes(32, 'big')
    
    return private_key, x_bytes + y_bytes  # 64 bytes total


def compute_ecdh_shared_key(private_key: ec.EllipticCurvePrivateKey, peer_pubkey_bytes: bytes) -> bytes:
    """
    Compute a shared secret using ECDH.
    
    Args:
        private_key: Local ECDH private key
        peer_pubkey_bytes: Peer's ephemeral public key (64 bytes, raw x||y format)
        
    Returns:
        Shared secret bytes
        
    Raises:
        ValueError: If peer_pubkey_bytes is not 64 bytes
    """
    if len(peer_pubkey_bytes) != 64:
        raise ValueError("Peer public key must be 64 bytes (x||y format)")
        
    x_int = int.from_bytes(peer_pubkey_bytes[:32], 'big')
    y_int = int.from_bytes(peer_pubkey_bytes[32:], 'big')
    
    # Create public key object from raw components
    peer_public_numbers = ec.EllipticCurvePublicNumbers(x_int, y_int, ec.SECP256R1())
    peer_public_key = peer_public_numbers.public_key()
    
    # Perform key exchange
    shared_key = private_key.exchange(ec.ECDH(), peer_public_key)
    return shared_key


def der_to_raw(der_sig: bytes) -> bytes:
    """
    Convert a DER-encoded ECDSA signature to a raw 64-byte signature (r||s).
    
    Args:
        der_sig: DER-encoded signature
        
    Returns:
        Raw 64-byte signature (r||s format), with each component padded to 32 bytes
    """
    r, s = decode_dss_signature(der_sig)
    r_bytes = r.to_bytes(32, byteorder='big')
    s_bytes = s.to_bytes(32, byteorder='big')
    return r_bytes + s_bytes


def raw_signature_to_der(raw_sig: bytes) -> bytes:
    """
    Convert a raw signature (64 bytes, concatenated r||s) to DER-encoded signature.
    
    Args:
        raw_sig: Raw 64-byte signature (r||s format)
        
    Returns:
        DER-encoded signature
        
    Raises:
        ValueError: If raw_sig is not 64 bytes
    """
    if len(raw_sig) != 64:
        raise ValueError("Raw signature must be 64 bytes (r||s).")
        
    r = int.from_bytes(raw_sig[:32], 'big')
    s = int.from_bytes(raw_sig[32:], 'big')
    return encode_dss_signature(r, s)
