"""Общий ML-пайплайн для провайдеров репутации (начало с Kaspersky)."""

from __future__ import annotations

import argparse
import asyncio
import logging
from pathlib import Path
from typing import Any

import pandas as pd
from catboost import CatBoostClassifier, Pool
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import sessionmaker

from ioc_checker.database import Cache, DatasetEntry, SessionLocal
from ioc_checker.providers.http_observer import HTTPObserverProvider
from ioc_checker.providers.otx import AlienVaultOTXProvider
from ioc_checker.providers.kaspersky import KasperskyProvider
from ioc_checker.providers.openpagerank import OpenPageRankProvider
from ioc_checker.providers.systemdns import SystemDNSProvider
from ioc_checker.providers.top1m import Top1MProvider
from ioc_checker.providers.whois import WhoisProvider
from ml.utils import (
    DecisionThreshold,
    ensure_directory,
    export_dataframe_stats,
    find_optimal_thresholds,
    log_feature_importance,
    log_model_metrics,
    parse_raw_payload,
    select_features_by_importance,
    write_training_report,
)

logger = logging.getLogger(__name__)


class GeneralMLPipeline:
    """Строит датафрейм фичей и обучает CatBoost для общих провайдеров."""

    def __init__(
        self,
        *,
        session_factory: sessionmaker[AsyncSession] = SessionLocal,
        model_dir: Path = Path("ml/models/general"),
        iterations: int = 1000,
        learning_rate: float = 0.1,
        depth: int = 6,
        early_stopping_rounds: int = 50,
        importance_percentile: float | None = 10.0,
        max_features: int | None = None,
    ) -> None:
        self._providers = [
            KasperskyProvider(),
            OpenPageRankProvider(),
            SystemDNSProvider(),
            Top1MProvider(),
            WhoisProvider(),
            HTTPObserverProvider(),
            AlienVaultOTXProvider(),
        ]
        self._session_factory = session_factory
        self._model_dir = model_dir
        self._label_column = "label"
        self._iterations = iterations
        self._learning_rate = learning_rate
        self._depth = depth
        self._early_stopping_rounds = early_stopping_rounds
        self._importance_percentile = importance_percentile
        self._max_features = max_features
        ensure_directory(self._model_dir)

    async def build_dataset_from_db(self) -> pd.DataFrame:
        """Собирает датафрейм фичей по доступным провайдерам из базы данных."""

        provider_names = [provider.name for provider in self._providers]
        stmt = (
            select(
                DatasetEntry.ioc,
                DatasetEntry.label,
                Cache.provider,
                Cache.raw_response,
            )
            .join(Cache, Cache.ioc == DatasetEntry.ioc)
            .where(Cache.provider.in_(provider_names))
        )

        async with self._session_factory() as session:
            logger.info(
                "Сбор записей dataset+cache для общих провайдеров: %s", provider_names
            )
            result = await session.execute(stmt)
            fetched_rows = result.fetchall()
            logger.info("Получено %d записей из ORM", len(fetched_rows))

        grouped: dict[str, dict[str, Any]] = {}
        for ioc, label, provider, raw_response in fetched_rows:
            entry = grouped.setdefault(ioc, {"label": label, "payloads": {}})
            entry["payloads"][provider] = parse_raw_payload(raw_response)

        dataset_rows: list[dict[str, Any]] = []
        for index, (ioc, entry) in enumerate(grouped.items()):
            payloads = entry.get("payloads", {})
            feature_row: dict[str, Any] = {}

            for provider in self._providers:
                provider_payload = payloads.get(provider.name, {})
                try:
                    provider_features = provider.build_features(
                        provider_payload, indicator=ioc
                    )
                except Exception:
                    logger.exception(
                        "Не удалось построить фичи провайдера %s для IOC %s",
                        provider.name,
                        ioc,
                    )
                    provider_features = {}
                feature_row.update(provider_features)

            feature_row[self._label_column] = int(entry.get("label"))

            if index < 3:
                logger.debug(
                    "Пример фичей для IOC %s: %s",
                    ioc,
                    list(feature_row.items())[:5],
                )
            dataset_rows.append(feature_row)

        dataframe = pd.DataFrame(dataset_rows)
        stats_path = self._model_dir / "general_dataset_stats.txt"
        export_dataframe_stats(dataframe, stats_path)
        logger.info(
            "Датафрейм собран: строк=%d, признаков=%d",
            len(dataframe),
            dataframe.shape[1],
        )
        return dataframe

    def _build_pool(self, split_df: pd.DataFrame, split_cat_columns: list[str]) -> Pool:
        split_features = split_df.drop(columns=[self._label_column])
        split_target = split_df[self._label_column]
        logger.debug(
            "Сплит: строк=%d, признаков=%d",
            len(split_features),
            split_features.shape[1],
        )
        return Pool(split_features, label=split_target, cat_features=split_cat_columns)

    def _fit_model(
        self,
        train_split: pd.DataFrame,
        test_split: pd.DataFrame,
        split_cat_columns: list[str],
    ) -> tuple[
        CatBoostClassifier,
        Pool,
        dict[str, float],
        list[tuple[str, float]],
        list[DecisionThreshold],
    ]:
        train_pool = self._build_pool(train_split, split_cat_columns)
        test_pool = self._build_pool(test_split, split_cat_columns)
        model = CatBoostClassifier(
            iterations=self._iterations,
            learning_rate=self._learning_rate,
            depth=self._depth,
            loss_function="Logloss",
            eval_metric="AUC",
            od_type="Iter",
            od_wait=self._early_stopping_rounds,
            use_best_model=True,
            verbose=100,
        )
        logger.info(
            "Запускаем обучение CatBoost: обучающих строк=%d, тестовых строк=%d, признаков=%d",
            len(train_split),
            len(test_split),
            train_split.drop(columns=[self._label_column]).shape[1],
        )
        model.fit(train_pool, eval_set=test_pool)

        metrics = log_model_metrics(model, test_pool)
        importance = log_feature_importance(
            model, train_split.drop(columns=[self._label_column]).columns.tolist()
        )
        probabilities = model.predict_proba(test_pool)[:, 1]
        thresholds = find_optimal_thresholds(
            истинные_метки=test_split[self._label_column].to_numpy(),
            вероятности_позитивного_класса=probabilities,
        )
        return model, test_pool, metrics, importance, thresholds

    def train_model(self, dataframe: pd.DataFrame) -> Path:
        """Обучает CatBoost модель, затем дообучает с отсечёнными признаками."""

        if len(dataframe) < 2:
            raise ValueError("Для обучения требуется минимум две записи")

        train_df = dataframe.sample(frac=0.8, random_state=42)
        test_df = dataframe.drop(train_df.index)
        if test_df.empty:
            test_df = train_df.iloc[:1]
            train_df = train_df.iloc[1:]
            logger.warning(
                "Тестовый сплит пустой, выделяем одну запись из обучающего набора"
            )

        features_df = dataframe.drop(columns=[self._label_column])
        cat_columns = [
            column
            for column in features_df.columns
            if features_df[column].dtype == "object"
        ]
        logger.debug("Категориальные признаки (полный набор): %s", cat_columns)

        (
            model_full,
            _,
            metrics_full,
            importance_full,
            thresholds_full,
        ) = self._fit_model(train_df, test_df, cat_columns)

        selected_features, applied_threshold = select_features_by_importance(
            importance_full,
            percentile=self._importance_percentile,
            max_features=self._max_features,
        )
        if not selected_features:
            logger.warning(
                "После отбора признаков ни один не прошёл порог %.4f — сохраняем полную модель",
                applied_threshold,
            )
            selected_features = features_df.columns.tolist()

        if len(selected_features) < len(features_df.columns):
            logger.info(
                "Отбираем %d/%d признаков по важности (порог=%.4f, максимум=%s)",
                len(selected_features),
                len(features_df.columns),
                applied_threshold,
                self._max_features or "∞",
            )
        else:
            logger.info(
                "Порог важности не сократил набор признаков, используем исходный"
            )

        pruned_train_df = train_df[selected_features + [self._label_column]]
        pruned_test_df = test_df[selected_features + [self._label_column]]
        pruned_cat_columns = [
            column
            for column in selected_features
            if dataframe[column].dtype == "object"
        ]

        pruned_stats_path = self._model_dir / "general_dataset_pruned_stats.txt"
        export_dataframe_stats(
            dataframe[selected_features + [self._label_column]], pruned_stats_path
        )
        logger.info(
            "Статистики по датасету с урезанными признаками сохранены в %s",
            pruned_stats_path,
        )

        (
            model_pruned,
            _,
            metrics_pruned,
            importance_pruned,
            thresholds_pruned,
        ) = self._fit_model(pruned_train_df, pruned_test_df, pruned_cat_columns)

        report_path = self._model_dir / "general_training_report.txt"
        write_training_report(
            report_path,
            metrics_full,
            importance_full,
            pruned_metrics=metrics_pruned,
            pruned_feature_importance=importance_pruned,
            thresholds=thresholds_full,
            pruned_thresholds=thresholds_pruned,
        )

        full_model_path = self._model_dir / "general_catboost_full.cbm"
        model_full.save_model(full_model_path)
        logger.info("Полная CatBoost модель сохранена в %s", full_model_path)

        pruned_model_path = self._model_dir / "general_catboost.cbm"
        model_pruned.save_model(pruned_model_path)
        logger.info(
            "CatBoost модель с отсечёнными признаками сохранена в %s", pruned_model_path
        )
        return pruned_model_path


async def _build_from_db(
    model_dir: Path,
    save_csv: Path | None,
    importance_percentile: float | None,
    max_features: int | None,
) -> None:
    """CLI: строит датасет из базы и обучает общую модель."""

    pipeline = GeneralMLPipeline(
        model_dir=model_dir,
        importance_percentile=importance_percentile,
        max_features=max_features,
    )
    dataset = await pipeline.build_dataset_from_db()
    if save_csv:
        ensure_directory(save_csv.parent)
        dataset.to_csv(save_csv, index=False)
        logger.info("Датасет с фичами сохранён в %s", save_csv)
    pipeline.train_model(dataset)


async def _main() -> None:
    parser = argparse.ArgumentParser(
        description="ML-пайплайн обучения для общих провайдеров",
    )
    parser.add_argument(
        "--model-dir",
        type=Path,
        default=Path("ml/models/general"),
        help="Директория для сохранения модели",
    )
    parser.add_argument(
        "--save-csv",
        type=Path,
        help="Путь для сохранения собранного датасета с фичами в CSV",
    )
    parser.add_argument(
        "--importance-percentile",
        type=float,
        default=10.0,
        help="Динамический порог важности по процентилю",
    )
    parser.add_argument(
        "--max-features",
        type=int,
        help="Максимальное количество оставляемых признаков (None — без ограничения)",
    )
    args = parser.parse_args()

    await _build_from_db(
        args.model_dir,
        args.save_csv,
        args.importance_percentile,
        args.max_features,
    )


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO, force=True)
    asyncio.run(_main())
