import base64
import time
from typing import Dict, Optional
from cryptography.hazmat.primitives.asymmetric import ed25519
from cryptography.hazmat.primitives import serialization

# Protocol Constants
SIGNATURE_ALGORITHM = 'ed25519'
MAX_TIME_SKEW_SECONDS = 30

class ValleyDamSigner:
    """Handles the cryptographic signing of HTTP requests."""

    def __init__(self, private_key_path: str):
        """
        Load an Ed25519 private key from a PEM file.
        
        Args:
            private_key_path: Path to the .pem file generated by the CLI.
        """
        try:
            with open(private_key_path, 'rb') as key_file:
                self.private_key = serialization.load_pem_private_key(
                    key_file.read(),
                    password=None
                )
        except FileNotFoundError:
            raise ValueError(f"Private key not found at: {private_key_path}")
        except ValueError:
            raise ValueError("Invalid PEM file format.")

    def sign(self, method: str, host: str, path: str) -> Dict[str, str]:
        """
        Generate authentication headers for a request.
        
        The payload format is: "{METHOD} {HOST} {PATH} {TIMESTAMP}"
        """
        timestamp = int(time.time())
        
        # Canonicalization: Create the immutable string to sign
        payload = f"{method.upper()} {host} {path} {timestamp}".encode('utf-8')
        
        # Sign using Ed25519
        sig_bytes = self.private_key.sign(payload)
        sig_b64 = base64.b64encode(sig_bytes).decode('utf-8')
        
        return {
            'X-ValleyDam-Signature': sig_b64,
            'X-ValleyDam-Time': str(timestamp),
            'X-ValleyDam-Method': SIGNATURE_ALGORITHM
        }


class ValleyDamVerifier:
    """Handles the verification of incoming request signatures."""

    def __init__(self, key_resolver_func):
        """
        Args:
            key_resolver_func: A callable that takes a key_id string and 
                               returns an Ed25519PublicKey object (or None).
        """
        self.resolve_key = key_resolver_func

    def verify(self, method: str, host: str, path: str, headers: Dict) -> bool:
        """
        Verify that a request was signed by the owner of the DNS record.

        Raises:
            ValueError: If headers are missing, timestamp is expired, or signature is invalid.
        """
        # 1. Extract Headers
        sig_b64 = headers.get('X-ValleyDam-Signature')
        timestamp = headers.get('X-ValleyDam-Time')
        key_id = headers.get('X-ValleyDam-KeyID') 

        if not all([sig_b64, timestamp, key_id]):
            raise ValueError("Missing required ValleyDam authentication headers.")

        # 2. Check Timestamp (Replay Attack Prevention)
        try:
            req_time = int(timestamp)
            now = int(time.time())
            if abs(now - req_time) > MAX_TIME_SKEW_SECONDS:
                raise ValueError(f"Request expired. Server time: {now}, Request time: {req_time}")
        except ValueError:
            raise ValueError("Invalid timestamp format.")

        # 3. Resolve Public Key
        public_key = self.resolve_key(key_id)
        if not public_key:
            raise ValueError(f"Public key not found for identity: {key_id}")

        # 4. Reconstruct and Verify Payload
        payload = f"{method.upper()} {host} {path} {timestamp}".encode('utf-8')

        try:
            sig_bytes = base64.b64decode(sig_b64)
            public_key.verify(sig_bytes, payload)
            return True
        except Exception:
            raise ValueError("Cryptographic verification failed. Signature does not match payload.")