from __future__ import annotations

import ipaddress
from dataclasses import dataclass, field
from typing import Dict, Optional

from app.core.logger import make_log


@dataclass
class ASNResolver:
    cache: Dict[str, int] = field(default_factory=dict)

    def normalise(self, ip: str | None) -> Optional[str]:
        if not ip:
            return None
        try:
            return str(ipaddress.ip_address(ip))
        except Exception:
            return None

    def resolve(self, ip: str | None) -> Optional[int]:
        norm = self.normalise(ip)
        if not norm:
            return None
        return self.cache.get(norm)

    def learn(self, ip: str, asn: int) -> None:
        norm = self.normalise(ip)
        if not norm:
            make_log("ASNResolver", f"Invalid IP provided for learn: {ip}", level="warning")
            return
        self.cache[norm] = asn


resolver = ASNResolver()