"""Сбор топовых HTTP-заголовков из кэша в базе данных."""

from __future__ import annotations

import argparse
import asyncio
import json
import logging
from pathlib import Path
from typing import Sequence

from sqlalchemy import and_, func, select
from sqlalchemy.dialects.postgresql import JSONB
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.sql import cast

from ioc_checker.database import Cache, DatasetEntry, SessionLocal
from ioc_checker.providers.http_observer import HTTPObserverProvider
from ml.utils import ensure_directory

logger = logging.getLogger(__name__)


async def fetch_top_headers(
    session: AsyncSession, provider: str, limit: int = 100
) -> list[str]:
    """Извлечь самые частые HTTP-заголовки провайдера."""

    dataset_filtered_cache = (
        select(Cache.raw_response.label("raw_response"))
        .where(Cache.provider == provider)
        .join(DatasetEntry, DatasetEntry.ioc == Cache.ioc)
        .subquery()
    )

    response_count = await session.scalar(
        select(func.count()).select_from(dataset_filtered_cache)
    )
    if not response_count:
        logger.warning(
            "Нет ответов в кэше для провайдера %s, пространство заголовков пустое",
            provider,
        )
        return []

    response_headers = cast(
        dataset_filtered_cache.c.raw_response["scan"]["response_headers"], JSONB
    )
    headers_subquery = (
        select(func.jsonb_object_keys(response_headers).label("header_key"))
        .select_from(dataset_filtered_cache)
        .subquery()
    )

    frequency_stmt = (
        select(
            headers_subquery.c.header_key,
            func.count().label("frequency"),
        )
        .group_by(headers_subquery.c.header_key)
        .having(
            and_(
                func.count() > 1,
                func.count() < response_count,
            )
        )
        .order_by(func.count().desc())
        .limit(limit)
    )

    result = await session.execute(frequency_stmt)
    rows = result.fetchall()
    headers = [str(row.header_key) for row in rows if row.header_key]
    logger.info(
        "Выбрано %d заголовков для провайдера %s из %d ответов",
        len(headers),
        provider,
        response_count,
    )
    logger.debug("Топ-5 заголовков: %s", headers[:5])
    return headers


def save_headers(headers: Sequence[str], output_path: Path) -> None:
    """Сохранить заголовки в JSON файл."""

    ensure_directory(output_path.parent)
    output_path.write_text(
        json.dumps(list(headers), ensure_ascii=False, indent=2), encoding="utf-8"
    )
    logger.info("Топ заголовков сохранён в %s", output_path)


async def _main_async(args: argparse.Namespace) -> None:
    provider = args.provider or HTTPObserverProvider().name
    logger.info(
        "Собираем топ %d заголовков для провайдера %s", args.limit, provider
    )
    async with SessionLocal() as session:
        headers = await fetch_top_headers(session, provider, args.limit)
    save_headers(headers, args.output)


def _parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Выгрузка топовых HTTP-заголовков из базы данных",
    )
    parser.add_argument(
        "--output",
        type=Path,
        default=Path("ioc_checker/providers/http_observer_top_headers.json"),
        help="Путь для сохранения топа заголовков",
    )
    parser.add_argument(
        "--provider",
        type=str,
        default=HTTPObserverProvider().name,
        help="Имя провайдера в базе данных",
    )
    parser.add_argument(
        "--limit",
        type=int,
        default=100,
        help="Количество заголовков для выгрузки",
    )
    return parser.parse_args()


def main() -> None:
    logging.basicConfig(level=logging.INFO, force=True)
    args = _parse_args()
    asyncio.run(_main_async(args))


if __name__ == "__main__":
    main()
