"""DoH-провайдер DNS-запросов Google."""

from __future__ import annotations

import asyncio
from contextlib import AbstractAsyncContextManager, asynccontextmanager
import logging
from typing import Any, AsyncIterator

import httpx

from ..iocs import IOCType
from .base import BaseProvider, FeatureValue

logger = logging.getLogger(__name__)


RECORD_TYPES: tuple[str, ...] = (
    "A",
    "AAAA",
    "MX",
    "NS",
    "TXT",
    "CNAME",
    "SOA",
    "CAA",
)


class SystemDNSProvider(BaseProvider):
    """Выполняет DNS-запросы через DoH Google."""

    DOH_BASE_URL = "https://dns.google"
    DOH_ENDPOINT = "/resolve"

    def __init__(
        self,
        doh_base_url: str | None = None,
        *,
        transport: httpx.BaseTransport | None = None,
    ) -> None:
        super().__init__("systemdns", requires_token=False)
        self.doh_base_url = doh_base_url or self.DOH_BASE_URL
        self._transport = transport

    def get_context(
        self, token: str | None = None
    ) -> AbstractAsyncContextManager[httpx.AsyncClient]:
        """Вернуть контекст с HTTP-клиентом DoH Google."""

        @asynccontextmanager
        async def _client() -> AsyncIterator[httpx.AsyncClient]:
            logger.debug(
                "Создаём HTTP-клиент для DoH: базовый URL %s, таймаут %.2f",
                self.doh_base_url,
                self.http_timeout,
            )
            client_kwargs = self.build_http_client_kwargs(
                base_url=self.doh_base_url,
                headers={
                    "Accept": "application/dns-json",
                    "User-Agent": "ioc-checker/systemdns",
                },
            )
            if self._transport is not None:
                client_kwargs["transport"] = self._transport
            async with httpx.AsyncClient(**client_kwargs) as client:
                yield client

        return _client()

    async def fetch(
        self, indicator: IOCType, context: Any
    ) -> dict[str, list[str]] | bool:
        """Выполнить DNS-запросы через DoH и вернуть агрегированные записи."""

        if indicator.category != "domain":
            logger.info(
                "SystemDNS поддерживает только домены, пропускаем %s (%s)",
                indicator.value,
                indicator.category,
            )
            return False

        logger.info("Запрашиваем DNS-записи через DoH для домена %s", indicator.value)

        client: httpx.AsyncClient | None = None
        if isinstance(context, httpx.AsyncClient):
            client = context

        async def _fetch_with_client(http_client: httpx.AsyncClient) -> dict[str, list[str]] | bool:
            try:
                tasks = [
                    self._request_records(http_client, indicator.value, record_type)
                    for record_type in RECORD_TYPES
                ]
                resolved_records = await asyncio.gather(*tasks)
                return {record_type: records for record_type, records in zip(RECORD_TYPES, resolved_records)}
            except Exception as exc:  # noqa: BLE001
                logger.error(
                    "Не удалось получить DoH-ответ для домена %s: %s",
                    indicator.value,
                    exc,
                )
                return False

        if client is not None:
            return await _fetch_with_client(client)

        async with self.get_context() as managed_client:
            return await _fetch_with_client(managed_client)

    async def _request_records(
        self, client: httpx.AsyncClient, domain: str, record_type: str
    ) -> list[str]:
        response = await client.get(
            self.DOH_ENDPOINT,
            params={"name": domain, "type": record_type},
        )
        response.raise_for_status()

        payload = response.json()
        answers = payload.get("Answer") or []
        records = [entry.get("data") for entry in answers if entry.get("data")]
        logger.debug(
            "Получены записи %s для %s: %s",
            record_type,
            domain,
            records,
        )
        return records

    def build_features(
        self,
        payload: Any,
        *,
        indicator: IOCType | None = None,
    ) -> dict[str, FeatureValue]:
        """Подготовить фичи на основе DNS-ответов."""

        record_counts: dict[str, int] = {record_type: 0 for record_type in RECORD_TYPES}
        for record_type in RECORD_TYPES:
            records = payload.get(record_type, [])
            record_counts[record_type] = len(records)

        txt_records = payload.get("TXT", [])

        dns_has_spf = any("v=spf1" in txt for txt in txt_records)
        dns_has_dmarc = any("v=dmarc1" in txt for txt in txt_records)
        dns_txt_length_total = sum(len(text) for text in txt_records)

        dns_complexity = sum(record_counts.values())

        features: dict[str, FeatureValue] = {
            "dns_a_count": record_counts["A"],
            "dns_aaaa_count": record_counts["AAAA"],
            "dns_mx_count": record_counts["MX"],
            "dns_ns_count": record_counts["NS"],
            "dns_txt_count": record_counts["TXT"],
            "dns_cname_count": record_counts["CNAME"],
            "dns_soa_count": record_counts["SOA"],
            "dns_caa_count": record_counts["CAA"],
            "dns_complexity": dns_complexity,
            "dns_has_spf": dns_has_spf,
            "dns_has_dmarc": dns_has_dmarc,
            "dns_txt_length_total": dns_txt_length_total,
        }
        return features
