"""Enrichment helpers for BGP observation and anomaly pipelines."""

from __future__ import annotations

import ipaddress
from collections import defaultdict
from datetime import UTC, datetime
from typing import Any

from sqlalchemy import Integer, cast, select, text
from sqlalchemy.ext.asyncio import AsyncSession

from app.core.countries import get_country_centroid, normalize_country
from app.models.bgp_observation import BGPObservation
from app.models.collected_data import CollectedData


def _safe_int(value: Any) -> int | None:
    try:
        if value in (None, ""):
            return None
        return int(value)
    except (TypeError, ValueError):
        return None


def _parse_timestamp(value: Any) -> datetime:
    if isinstance(value, datetime):
        return value.astimezone(UTC) if value.tzinfo else value.replace(tzinfo=UTC)

    if isinstance(value, (int, float)):
        return datetime.fromtimestamp(value, tz=UTC)

    if isinstance(value, str) and value:
        normalized = value.replace("Z", "+00:00")
        parsed = datetime.fromisoformat(normalized)
        return parsed.astimezone(UTC) if parsed.tzinfo else parsed.replace(tzinfo=UTC)

    return datetime.now(UTC)


def _dedupe_as_path(as_path: list[int]) -> list[int]:
    deduped: list[int] = []
    for asn in as_path:
        if not deduped or deduped[-1] != asn:
            deduped.append(asn)
    return deduped


def _compact_locations(items: list[dict[str, Any]]) -> list[dict[str, Any]]:
    results: list[dict[str, Any]] = []
    seen: set[tuple[Any, ...]] = set()
    for item in items:
        key = (
            item.get("country"),
            item.get("city"),
            item.get("latitude"),
            item.get("longitude"),
        )
        if key in seen:
            continue
        seen.add(key)
        results.append(item)
    return results


def extract_bgp_network_fields(prefix: str) -> dict[str, Any]:
    if not prefix:
        return {
            "prefix_family": None,
            "prefix_length": None,
            "prefix_supernet": None,
            "is_more_specific": False,
        }

    try:
        network = ipaddress.ip_network(prefix, strict=False)
    except ValueError:
        return {
            "prefix_family": None,
            "prefix_length": None,
            "prefix_supernet": None,
            "is_more_specific": False,
        }

    supernet_prefix = 16 if network.version == 4 else 32
    if network.prefixlen > supernet_prefix:
        prefix_supernet = str(network.supernet(new_prefix=supernet_prefix))
    else:
        prefix_supernet = str(network)

    return {
        "prefix_family": f"ipv{network.version}",
        "prefix_length": int(network.prefixlen),
        "prefix_supernet": prefix_supernet,
        "is_more_specific": network.prefixlen > (24 if network.version == 4 else 48),
    }


async def _lookup_prefix_geography(
    db: AsyncSession,
    prefix_values: list[str],
) -> dict[str, dict[str, Any]]:
    async def _query_prefix_metadata(
        *,
        source: str,
        family: str,
        range_start: str,
        range_end: str,
    ) -> dict[str, Any] | None:
        result = await db.execute(
            text(
                """
                SELECT metadata
                FROM collected_data
                WHERE source = :source
                  AND COALESCE(is_current, TRUE) = TRUE
                  AND metadata->>'family' = :family
                  AND CAST(metadata->>'range_start' AS inet) <= CAST(:range_start AS inet)
                  AND CAST(metadata->>'range_end' AS inet) >= CAST(:range_end AS inet)
                ORDER BY
                  masklen(CAST(metadata->>'prefix' AS cidr)) DESC NULLS LAST,
                  id DESC
                LIMIT 1
                """
            ),
            {
                "source": source,
                "family": family,
                "range_start": range_start,
                "range_end": range_end,
            },
        )
        row = result.fetchone()
        if not row:
            return None

        if isinstance(row, dict):
            payload = row.get("metadata") or row.get("extra_data")
        elif hasattr(row, "_mapping"):
            payload = row._mapping.get("metadata") or row._mapping.get("extra_data")
        else:
            payload = row[0]

        return payload if isinstance(payload, dict) else None

    results: dict[str, dict[str, Any]] = {}

    for prefix in prefix_values:
        try:
            network = ipaddress.ip_network(prefix, strict=False)
        except ValueError:
            continue

        family = f"ipv{network.version}"
        range_start = str(network.network_address)
        range_end = str(network.broadcast_address)
        payload = await _query_prefix_metadata(
            source="opengeofeed_prefix_geo",
            family=family,
            range_start=range_start,
            range_end=range_end,
        )
        selected_source = "opengeofeed"
        if not payload:
            payload = await _query_prefix_metadata(
                source="iptoasn_prefix_geo",
                family=family,
                range_start=range_start,
                range_end=range_end,
            )
            selected_source = "iptoasn"
        if not payload:
            payload = await _query_prefix_metadata(
                source="nro_delegated_prefix_geo",
                family=family,
                range_start=range_start,
                range_end=range_end,
            )
            selected_source = "nro_delegated"
        if not payload:
            continue

        country = normalize_country(payload.get("country") or payload.get("country_code"))
        prefix_hint = payload.get("prefix") or prefix
        asn = _safe_int(payload.get("asn"))
        as_name = payload.get("as_name")
        city = payload.get("city")
        centroid = get_country_centroid(country)
        regions = []
        if country:
            regions.append(
                {
                    "country": country,
                    "city": city,
                    "latitude": centroid.get("latitude") if centroid else None,
                    "longitude": centroid.get("longitude") if centroid else None,
                }
            )

        results[prefix] = {
            "prefix": prefix_hint,
            "country": country,
            "city": city,
            "asn": asn,
            "as_name": as_name,
            "source": payload.get("source_dataset")
            or (
                "opengeofeed_public"
                if selected_source == "opengeofeed"
                else (
                    "iptoasn_combined"
                    if selected_source == "iptoasn"
                    else "nro_delegated_stats"
                )
            ),
            "confidence": payload.get("confidence")
            or (
                "geofeed"
                if selected_source == "opengeofeed"
                else (
                    "country_range"
                    if selected_source == "iptoasn"
                    else "registry_allocated"
                )
            ),
            "geography_mode": "prefix_geography",
            "regions": regions,
        }

    return results


async def lookup_prefix_geography(
    db: AsyncSession,
    prefix_values: list[str],
) -> dict[str, dict[str, Any]]:
    return await _lookup_prefix_geography(db, prefix_values)


async def enrich_bgp_events_for_batch(
    db: AsyncSession,
    *,
    source: str,
    events: list[dict[str, Any]],
) -> list[dict[str, Any]]:
    if not events:
        return []

    prefixes = {
        str((event.get("metadata") or {}).get("prefix") or "").strip()
        for event in events
        if (event.get("metadata") or {}).get("prefix")
    }
    prefix_values = sorted(prefix for prefix in prefixes if prefix)
    origin_asns = sorted(
        {
            asn
            for event in events
            for asn in [
                _safe_int((event.get("metadata") or {}).get("origin_asn")),
                _safe_int((event.get("metadata") or {}).get("new_origin_asn")),
            ]
            if asn is not None
        }
    )

    historical_prefix_baseline: dict[str, dict[str, Any]] = {}
    if prefix_values:
        previous_result = await db.execute(
            select(
                BGPObservation.prefix,
                BGPObservation.origin_asn,
                BGPObservation.collector,
                BGPObservation.collector_geo,
            ).where(
                BGPObservation.source == source,
                BGPObservation.prefix.in_(prefix_values),
            )
        )
        by_prefix: defaultdict[str, list[dict[str, Any]]] = defaultdict(list)
        for prefix, origin_asn, collector, collector_geo in previous_result.all():
            if prefix:
                by_prefix[str(prefix)].append(
                    {
                        "origin_asn": origin_asn,
                        "collector": collector,
                        "collector_geo": collector_geo or {},
                    }
                )

        for prefix, observations in by_prefix.items():
            unique_origins = sorted(
                {
                    observation["origin_asn"]
                    for observation in observations
                    if observation["origin_asn"] is not None
                }
            )
            unique_collectors = sorted(
                {
                    observation["collector"]
                    for observation in observations
                    if observation["collector"]
                }
            )
            historical_prefix_baseline[prefix] = {
                "historical_origin_asns": unique_origins,
                "historical_collectors": unique_collectors,
                "historical_observation_count": len(observations),
                "historical_regions": _compact_locations(
                    [
                        observation["collector_geo"] or {}
                        for observation in observations
                        if observation["collector_geo"]
                    ]
                ),
            }

    asn_profiles: dict[int, dict[str, Any]] = {}
    prefix_geographies = await _lookup_prefix_geography(db, prefix_values) if prefix_values else {}
    if origin_asns:
        peeringdb_result = await db.execute(
            select(CollectedData)
            .where(CollectedData.source == "peeringdb_network")
            .where(CollectedData.is_current.is_(True))
            .where(
                cast(CollectedData.extra_data["asn"].as_string(), Integer).in_(origin_asns),
            )
            .order_by(CollectedData.id.desc())
        )
        for record in peeringdb_result.scalars().all():
            metadata = record.extra_data or {}
            asn = _safe_int(metadata.get("asn"))
            if asn is None or asn not in origin_asns:
                continue
            current = asn_profiles.get(asn)
            if current and (current.get("id") or 0) > (record.id or 0):
                continue
            asn_profiles[asn] = {
                "id": record.id,
                "asn": asn,
                "name": record.name,
                "country": metadata.get("country"),
                "city": metadata.get("city"),
                "source": "peeringdb_network",
                "info_type": metadata.get("info_type"),
                "info_traffic": metadata.get("info_traffic"),
                "info_ratio": metadata.get("info_ratio"),
                "ix_count": metadata.get("ix_count"),
                "url": metadata.get("url"),
            }

    collector_counts: defaultdict[str, int] = defaultdict(int)
    for event in events:
        collector = (event.get("metadata") or {}).get("collector")
        if collector:
            collector_counts[str(collector)] += 1

    enriched: list[dict[str, Any]] = []
    for event in events:
        metadata = dict(event.get("metadata") or {})
        prefix = str(metadata.get("prefix") or "").strip()
        as_path = metadata.get("as_path") or []
        normalized_as_path = [asn for asn in (_safe_int(item) for item in as_path) if asn is not None]
        deduped_as_path = _dedupe_as_path(normalized_as_path)
        collector = str(metadata.get("collector") or "").strip()
        collector_location = metadata.get("collector_location") or {}
        baseline = historical_prefix_baseline.get(prefix, {})
        prefix_geography = prefix_geographies.get(prefix)
        observed_at = _parse_timestamp(metadata.get("timestamp") or event.get("reference_date"))
        origin_asn = _safe_int(metadata.get("origin_asn"))
        new_origin_asn = _safe_int(metadata.get("new_origin_asn"))
        baseline_regions = baseline.get("historical_regions", [])
        prefix_scope_regions = _compact_locations([*baseline_regions])

        enrichment = {
            **extract_bgp_network_fields(prefix),
            "observed_at": observed_at.isoformat(),
            "normalized_as_path": normalized_as_path,
            "deduped_as_path": deduped_as_path,
            "deduped_as_path_length": len(deduped_as_path),
            "path_prepending": len(normalized_as_path) > len(deduped_as_path),
            "collector_region": {
                "city": collector_location.get("city"),
                "country": collector_location.get("country"),
            },
            "collector_observation_count_in_batch": collector_counts.get(collector, 0),
            "batch_visibility_collectors": sorted(collector_counts.keys()),
            "prefix_baseline": baseline,
            "is_new_origin_for_prefix": (
                origin_asn is not None
                and origin_asn
                not in set(baseline.get("historical_origin_asns", []))
            ),
            "rpki_validation": {
                "status": "unknown",
                "reason": "no_rpki_roa_dataset_configured",
            },
            "origin_asn_profile": asn_profiles.get(origin_asn),
            "new_origin_asn_profile": asn_profiles.get(new_origin_asn),
            "prefix_geography": prefix_geography,
            "prefix_scope": {
                "countries": sorted(
                    {
                        item.get("country")
                        for item in prefix_scope_regions
                        if item.get("country")
                    }
                ),
                "cities": sorted(
                    {
                        item.get("city")
                        for item in prefix_scope_regions
                        if item.get("city")
                    }
                ),
                "regions": prefix_scope_regions,
            },
        }

        enriched.append(
            {
                **event,
                "metadata": {
                    **metadata,
                    "enrichment": enrichment,
                },
            }
        )

    return enriched