"""
IP datatype for VisiData.

Goals:
- One type (`ip`) that accepts IPv4/IPv6 addresses and CIDR networks (e.g. 192.168.7.0/24).
- Correct sorting (by version then numeric value, with networks after addresses).
- `ip * netmask` returns True/False for membership (address in network).
- Enrichment properties with normalized shapes:
    - ip.ipinfo.country
    - ip.asn.asn / ip.asn.name / ip.asn.country
    - ip.vt.verdict  (e.g. "3/94")

Network calls are optional; when deps/keys are missing, properties return empty defaults.
"""

from __future__ import annotations

from dataclasses import dataclass
import functools
import os
import pickle
import sqlite3
import time
from typing import Any, Dict, Optional, Tuple, Union

import ipaddress

from visidata import vd
from visidata.sheets import TableSheet


_CACHE_PATH = os.path.expanduser("~/.visidata_cache.db")


def _sqlite_getset(key: str, fn, *, max_age: Optional[int] = None):
    """Tiny sqlite+pickle cache. Falls back to computing if db can't be used."""
    try:
        os.makedirs(os.path.dirname(_CACHE_PATH), exist_ok=True)
        with sqlite3.connect(_CACHE_PATH) as conn:
            conn.execute(
                "CREATE TABLE IF NOT EXISTS cache (key TEXT PRIMARY KEY, value BLOB, timestamp INTEGER)"
            )
            cur = conn.cursor()
            cur.execute("SELECT value, timestamp FROM cache WHERE key=?", (key,))
            row = cur.fetchone()
            now = int(time.time())
            if row:
                val_blob, ts = row
                if max_age is None or now - int(ts) <= max_age:
                    return pickle.loads(val_blob)
            val = fn()
            cur.execute(
                "INSERT OR REPLACE INTO cache (key, value, timestamp) VALUES (?, ?, ?)",
                (key, pickle.dumps(val), now),
            )
            conn.commit()
            return val
    except Exception:
        return fn()


def _is_nullish(v: Any) -> bool:
    return v is None or v == "" or v == "null"


@dataclass(frozen=True)
class IPInfo:
    raw: Optional[Dict[str, Any]] = None

    @property
    def country(self) -> str:
        return str((self.raw or {}).get("country") or "")

    @property
    def region(self) -> str:
        return str((self.raw or {}).get("region") or "")

    @property
    def city(self) -> str:
        return str((self.raw or {}).get("city") or "")

    @property
    def org(self) -> str:
        return str((self.raw or {}).get("org") or "")

    @property
    def asn(self) -> str:
        # ipinfo "org" is often "AS12345 Some Org"
        org = self.org
        if org.startswith("AS"):
            tok = org.split(" ", 1)[0]
            return tok
        return ""


@dataclass(frozen=True)
class ASNInfo:
    asn: str = ""
    name: str = ""
    country: str = ""
    raw: Optional[Dict[str, Any]] = None


@dataclass(frozen=True)
class VTInfo:
    malicious: int = 0
    suspicious: int = 0
    total: int = 0
    reputation: Optional[int] = None
    categories: Tuple[str, ...] = ()
    raw: Optional[Dict[str, Any]] = None

    @property
    def verdict(self) -> str:
        if self.total <= 0:
            return ""
        return f"{self.malicious}/{self.total}"

    @property
    def ratio(self) -> Optional[float]:
        if self.total <= 0:
            return None
        return self.malicious / self.total

    # User asked for vt.type; "category" is a better name, but keep alias.
    @property
    def type(self) -> str:  # noqa: A003 - intentional API
        return self.category

    @property
    def category(self) -> str:
        return self.categories[0] if self.categories else ""


def _read_vt_key() -> str:
    try:
        with open(os.path.expanduser("~/.virustotal_api_key")) as f:
            return f.readline().strip()
    except Exception:
        return ""


@functools.lru_cache(maxsize=4096)
def _ipinfo_raw(ip: str) -> Optional[Dict[str, Any]]:
    def _do():
        try:
            import requests  # optional dep

            r = requests.get(f"http://ipinfo.io/{ip}/json", timeout=5)
            return r.json()
        except Exception:
            return None

    return _sqlite_getset(f"ipinfo:{ip}", _do, max_age=60 * 60 * 24)


@functools.lru_cache(maxsize=4096)
def _asn_info(ip: str) -> ASNInfo:
    # Prefer ipinfo-derived ASN (no scraping deps).
    raw = _ipinfo_raw(ip) or {}
    org = str(raw.get("org") or "")
    country = str(raw.get("country") or "")
    asn = ""
    name = ""
    if org.startswith("AS"):
        parts = org.split(" ", 1)
        asn = parts[0]
        name = parts[1] if len(parts) > 1 else ""
    return ASNInfo(asn=asn, name=name, country=country, raw={"org": org, "country": country})


@functools.lru_cache(maxsize=4096)
def _vt_info(ip: str) -> VTInfo:
    key = _read_vt_key()
    if not key:
        return VTInfo()

    def _do() -> VTInfo:
        try:
            import requests  # optional dep

            r = requests.get(
                f"https://www.virustotal.com/api/v3/ip_addresses/{ip}",
                headers={"x-apikey": key},
                timeout=10,
            )
            data = r.json() if r.ok else None
            attrs = (((data or {}).get("data") or {}).get("attributes") or {})
            stats = attrs.get("last_analysis_stats") or {}
            malicious = int(stats.get("malicious") or 0)
            suspicious = int(stats.get("suspicious") or 0)
            total = 0
            try:
                total = int(sum(int(v or 0) for v in stats.values()))
            except Exception:
                total = 0

            cats = attrs.get("categories") or {}
            if isinstance(cats, dict):
                # keep stable order
                categories = tuple(str(v) for v in cats.values() if v)
            elif isinstance(cats, list):
                categories = tuple(str(v) for v in cats if v)
            else:
                categories = ()

            rep = attrs.get("reputation")
            try:
                reputation = int(rep) if rep is not None else None
            except Exception:
                reputation = None

            return VTInfo(
                malicious=malicious,
                suspicious=suspicious,
                total=total,
                reputation=reputation,
                categories=categories,
                raw=data,
            )
        except Exception:
            return VTInfo()

    return _sqlite_getset(f"vt_ip:{ip}", _do, max_age=60 * 60 * 24)


@functools.total_ordering
class IPValue:
    """Represents either an address or a network."""

    __slots__ = ("_obj",)

    def __init__(self, obj: Union[ipaddress._BaseAddress, ipaddress._BaseNetwork]):
        self._obj = obj

    @property
    def obj(self):
        return self._obj

    @property
    def version(self) -> int:
        return int(self._obj.version)

    @property
    def is_network(self) -> bool:
        return isinstance(self._obj, ipaddress._BaseNetwork)

    @property
    def is_address(self) -> bool:
        return isinstance(self._obj, ipaddress._BaseAddress)

    @property
    def sort_key(self) -> Tuple[int, int, int, int]:
        # (version, kind, addrint, prefixlen)
        if self.is_network:
            n: ipaddress._BaseNetwork = self._obj  # type: ignore[assignment]
            return (self.version, 1, int(n.network_address), int(n.prefixlen))
        a: ipaddress._BaseAddress = self._obj  # type: ignore[assignment]
        return (self.version, 0, int(a), 0)

    def __str__(self) -> str:
        return str(self._obj)

    def __repr__(self) -> str:
        return f"IPValue({self._obj!r})"

    def __hash__(self) -> int:
        return hash(self.sort_key)

    def __eq__(self, other: object) -> bool:
        if not isinstance(other, IPValue):
            return False
        return self.sort_key == other.sort_key

    def __lt__(self, other: object) -> bool:
        if not isinstance(other, IPValue):
            return NotImplemented
        return self.sort_key < other.sort_key

    def _coerce_network(self, other: Any) -> Optional[ipaddress._BaseNetwork]:
        if other is None:
            return None
        if isinstance(other, IPValue) and other.is_network:
            return other.obj  # type: ignore[return-value]
        if isinstance(other, ipaddress._BaseNetwork):
            return other
        s = str(other).strip()
        if not s:
            return None
        try:
            return ipaddress.ip_network(s, strict=False)
        except Exception:
            return None

    def _coerce_address(self, other: Any) -> Optional[ipaddress._BaseAddress]:
        if other is None:
            return None
        if isinstance(other, IPValue) and other.is_address:
            return other.obj  # type: ignore[return-value]
        if isinstance(other, ipaddress._BaseAddress):
            return other
        s = str(other).strip()
        if not s:
            return None
        try:
            return ipaddress.ip_address(s)
        except Exception:
            return None

    def __mul__(self, other: Any):
        # "ip * netmask" -> membership test (address in network).
        net = self._coerce_network(other)
        if self.is_address and net is not None:
            try:
                return self.obj in net  # type: ignore[operator]
            except Exception:
                return False
        addr = self._coerce_address(other)
        if self.is_network and addr is not None:
            try:
                return addr in self.obj  # type: ignore[operator]
            except Exception:
                return False
        return False

    def __rmul__(self, other: Any):
        # allow "netmask * ip" too
        return self.__mul__(other)

    # Normalized enrichments
    @property
    def ipinfo(self) -> IPInfo:
        if not self.is_address:
            return IPInfo()
        return IPInfo(_ipinfo_raw(str(self)))

    @property
    def asn(self) -> ASNInfo:
        if not self.is_address:
            return ASNInfo()
        return _asn_info(str(self))

    @property
    def vt(self) -> VTInfo:
        if not self.is_address:
            return VTInfo()
        return _vt_info(str(self))


def ip(val: Any) -> Optional[IPValue]:
    """VisiData type converter: parse IPv4/IPv6 address or CIDR network."""
    if _is_nullish(val):
        return None
    if isinstance(val, IPValue):
        return val
    if isinstance(val, (ipaddress._BaseAddress, ipaddress._BaseNetwork)):
        return IPValue(val)
    s = str(val).strip()
    if not s:
        return None

    # Strip common "ip:port" for IPv4 only (keep IPv6 intact).
    if "/" not in s and s.count(":") == 1:
        host, port = s.rsplit(":", 1)
        if port.isdigit():
            s = host

    try:
        if "/" in s:
            return IPValue(ipaddress.ip_network(s, strict=False))
        return IPValue(ipaddress.ip_address(s))
    except Exception:
        return None


vd.addType(ip, icon=":", formatter=lambda fmt, v: "" if v is None else str(v), name="IP")

TableSheet.addCommand(
    None,
    "type-ip",
    "cursorCol.type=ip",
    "set type of current column to IP (IPv4/IPv6/CIDR)",
)

