import dns.resolver
import base64
from functools import lru_cache
from typing import Optional
from cryptography.hazmat.primitives.asymmetric import ed25519
from .core import ValleyDamVerifier

# --- CACHED LOOKUP FUNCTION ---
@lru_cache(maxsize=256)
def _fetch_dns_record(key_id: str) -> Optional[ed25519.Ed25519PublicKey]:
    """
    Fetch and cache the Ed25519 public key from a DNS TXT record.
    Format expected: v=vd1; k=ed25519; p=<BASE64_KEY>;
    """
    if not key_id.startswith("dns:"):
        return None
    
    try:
        domain = key_id.split(":", 1)[1]
    except IndexError:
        return None

    # Use Google and Cloudflare DNS to avoid local ISP caching issues
    resolver = dns.resolver.Resolver(configure=False)
    resolver.nameservers = ['8.8.8.8', '1.1.1.1']
    resolver.lifetime = 2.0  # Timeout after 2 seconds

    try:
        txt_target = f"_agent.{domain}"
        answers = resolver.resolve(txt_target, 'TXT')
        
        for rdata in answers:
            # Clean up the TXT record string
            txt_string = rdata.to_text().replace('"', '').strip()
            
            # Parse key-value pairs
            parts = {}
            for item in txt_string.split(';'):
                if '=' in item:
                    k, v = item.strip().split('=', 1)
                    parts[k] = v
            
            # Return the key if found
            if parts.get('p'):
                try:
                    pub_bytes = base64.b64decode(parts['p'])
                    return ed25519.Ed25519PublicKey.from_public_bytes(pub_bytes)
                except Exception:
                    continue  # Malformed key, try next record

    except (dns.resolver.NXDOMAIN, dns.resolver.NoAnswer, dns.exception.Timeout):
        return None
    except Exception as e:
        # In production, you might want to log this error
        # print(f"DNS Error: {e}") 
        return None
    
    return None


class DnsKeyResolver:
    """Wrapper class for the cached DNS lookup."""
    def resolve(self, key_id: str):
        return _fetch_dns_record(key_id)


# Singleton instance to maintain cache across requests
GLOBAL_RESOLVER = DnsKeyResolver()


def verify_request(request) -> bool:
    """
    Helper function for Flask/Django/FastAPI.
    
    Usage:
        try:
            verify_request(request)
        except ValueError as e:
            abort(403, str(e))
    """
    verifier = ValleyDamVerifier(GLOBAL_RESOLVER.resolve)
    
    # Handle Flask-style request objects
    host = request.host
    path = request.path
    
    return verifier.verify(request.method, host, path, request.headers)