"""Модуль аудита уязвимостей для моделей KokaoCore."""
import torch
import numpy as np
from typing import Optional, Tuple, List, Dict, Any, Callable
import logging
from dataclasses import dataclass, field
from pathlib import Path
import json
from datetime import datetime

from .core import KokaoCore
from .core_base import CoreConfig

logger = logging.getLogger(__name__)


@dataclass
class VulnerabilityReport:
    """Отчет об уязвимостях модели."""
    model_info: Dict[str, Any]
    audit_date: str
    vulnerabilities: List[Dict[str, Any]] = field(default_factory=list)
    risk_score: float = 0.0
    recommendations: List[str] = field(default_factory=list)

    def add_vulnerability(self, name: str, severity: str, 
                          description: str, details: Dict[str, Any]) -> None:
        """Добавление уязвимости в отчет."""
        self.vulnerabilities.append({
            'name': name,
            'severity': severity,  # 'low', 'medium', 'high', 'critical'
            'description': description,
            'details': details
        })

        # Пересчет риска
        severity_weights = {'low': 1, 'medium': 2, 'high': 3, 'critical': 4}
        self.risk_score = sum(
            severity_weights.get(v['severity'], 0) 
            for v in self.vulnerabilities
        ) / max(len(self.vulnerabilities), 1) * 25

    def add_recommendation(self, recommendation: str) -> None:
        """Добавление рекомендации."""
        self.recommendations.append(recommendation)

    def save(self, path: str) -> None:
        """Сохранение отчета."""
        report_dict = {
            'model_info': self.model_info,
            'audit_date': self.audit_date,
            'vulnerabilities': self.vulnerabilities,
            'risk_score': self.risk_score,
            'recommendations': self.recommendations
        }
        Path(path).parent.mkdir(parents=True, exist_ok=True)
        with open(path, 'w') as f:
            json.dump(report_dict, f, indent=2)

    @classmethod
    def load(cls, path: str) -> 'VulnerabilityReport':
        """Загрузка отчета."""
        with open(path) as f:
            report_dict = json.load(f)
        return cls(**report_dict)


class VulnerabilityAuditor:
    """
    Аудитор уязвимостей для моделей KokaoCore.
    """

    def __init__(self, core: KokaoCore):
        """
        Инициализация аудитора.

        Args:
            core: Модель для аудита
        """
        self.core = core
        self.report: Optional[VulnerabilityReport] = None

    def start_audit(self) -> VulnerabilityReport:
        """Начало аудита."""
        self.report = VulnerabilityReport(
            model_info={
                'input_dim': self.core.config.input_dim,
                'device': str(self.core.device),
                'version': self.core.version,
                'is_quantized': self.core.is_quantized
            },
            audit_date=datetime.now().isoformat()
        )
        return self.report

    def check_weight_anomalies(self, threshold: float = 10.0) -> VulnerabilityReport:
        """
        Проверка аномалий в весах.

        Args:
            threshold: Порог для обнаружения аномалий

        Returns:
            Отчет об уязвимостях
        """
        if self.report is None:
            self.start_audit()

        with torch.no_grad():
            eff_w_plus, eff_w_minus = self.core._get_effective_weights()

            # Проверка на очень большие веса
            max_w_plus = eff_w_plus.max().item()
            max_w_minus = eff_w_minus.max().item()

            if max_w_plus > threshold or max_w_minus > threshold:
                self.report.add_vulnerability(
                    name='Large Weights',
                    severity='medium',
                    description='Обнаружены аномально большие веса',
                    details={
                        'max_w_plus': max_w_plus,
                        'max_w_minus': max_w_minus,
                        'threshold': threshold
                    }
                )
                self.report.add_recommendation(
                    'Примените регуляризацию или нормализацию весов'
                )

            # Проверка на нулевые веса
            zero_ratio_plus = (eff_w_plus < 1e-10).float().mean().item()
            zero_ratio_minus = (eff_w_minus < 1e-10).float().mean().item()

            if zero_ratio_plus > 0.5 or zero_ratio_minus > 0.5:
                self.report.add_vulnerability(
                    name='Dead Neurons',
                    severity='high',
                    description='Более 50% весов близки к нулю',
                    details={
                        'zero_ratio_plus': zero_ratio_plus,
                        'zero_ratio_minus': zero_ratio_minus
                    }
                )
                self.report.add_recommendation(
                    'Инициализируйте веса заново или уменьшите скорость обучения'
                )

        return self.report

    def check_gradient_sensitivity(self, num_samples: int = 100,
                                    perturbation: float = 0.01) -> VulnerabilityReport:
        """
        Проверка чувствительности к градиентам.

        Args:
            num_samples: Количество тестовых сэмплов
            perturbation: Размер возмущения

        Returns:
            Отчет об уязвимостях
        """
        if self.report is None:
            self.start_audit()

        sensitivities = []

        for _ in range(num_samples):
            x = torch.randn(self.core.config.input_dim)

            # Оригинальный сигнал
            s_original = self.core.signal(x)

            # Возмущенный вход
            x_perturbed = x + torch.randn_like(x) * perturbation
            s_perturbed = self.core.signal(x_perturbed)

            # Чувствительность
            sensitivity = abs(s_original - s_perturbed) / perturbation
            sensitivities.append(sensitivity.item())

        avg_sensitivity = np.mean(sensitivities)
        max_sensitivity = np.max(sensitivities)

        if max_sensitivity > 100:
            self.report.add_vulnerability(
                name='High Gradient Sensitivity',
                severity='high',
                description='Модель чрезмерно чувствительна к малым возмущениям',
                details={
                    'avg_sensitivity': avg_sensitivity,
                    'max_sensitivity': max_sensitivity
                }
            )
            self.report.add_recommendation(
                'Добавьте шум к данным обучения или используйте регуляризацию'
            )

        return self.report

    def check_adversarial_vulnerability(self, num_samples: int = 50,
                                         epsilon: float = 0.1,
                                         target_signal: float = 0.0) -> VulnerabilityReport:
        """
        Проверка уязвимости к состязательным атакам.

        Args:
            num_samples: Количество тестовых сэмплов
            epsilon: Максимальное возмущение
            target_signal: Целевой сигнал атаки

        Returns:
            Отчет об уязвимостях
        """
        if self.report is None:
            self.start_audit()

        success_rate = 0.0

        for _ in range(num_samples):
            x = torch.randn(self.core.config.input_dim)
            original_signal = self.core.signal(x)

            # FGSM-подобная атака
            x.requires_grad_(True)
            signal = self.core.forward(x)
            loss = (signal - target_signal) ** 2
            loss.backward()

            with torch.no_grad():
                x_adv = x + epsilon * torch.sign(x.grad)

            adv_signal = self.core.signal(x_adv)

            # Проверка успеха атаки
            if abs(adv_signal - target_signal) < abs(original_signal - target_signal):
                success_rate += 1

        success_rate /= num_samples

        if success_rate > 0.8:
            self.report.add_vulnerability(
                name='Adversarial Vulnerability',
                severity='critical',
                description='Модель уязвима к состязательным атакам',
                details={
                    'attack_success_rate': success_rate,
                    'epsilon': epsilon
                }
            )
            self.report.add_recommendation(
                'Используйте состязательное обучение для повышения устойчивости'
            )

        return self.report

    def check_numerical_stability(self) -> VulnerabilityReport:
        """
        Проверка численной стабильности.

        Returns:
            Отчет об уязвимостях
        """
        if self.report is None:
            self.start_audit()

        issues = []

        # Проверка на NaN/Inf в весах
        with torch.no_grad():
            if torch.isnan(self.core.w_plus).any():
                issues.append('NaN в w_plus')
            if torch.isinf(self.core.w_plus).any():
                issues.append('Inf в w_plus')
            if torch.isnan(self.core.w_minus).any():
                issues.append('NaN в w_minus')
            if torch.isinf(self.core.w_minus).any():
                issues.append('Inf в w_minus')

        if issues:
            self.report.add_vulnerability(
                name='Numerical Instability',
                severity='critical',
                description='Обнаружены числовые нестабильности',
                details={'issues': issues}
            )
            self.report.add_recommendation(
                'Проверьте входные данные и используйте стабильные функции активации'
            )

        # Проверка очень малых значений
        eff_w_plus, eff_w_minus = self.core._get_effective_weights()
        if (eff_w_plus < 1e-30).any():
            issues.append('Очень малые значения в w_plus')
        if (eff_w_minus < 1e-30).any():
            issues.append('Очень малые значения в w_minus')

        return self.report

    def check_overfitting(self, train_data: List[Tuple[torch.Tensor, float]],
                          test_data: List[Tuple[torch.Tensor, float]]) -> VulnerabilityReport:
        """
        Проверка переобучения.

        Args:
            train_data: Обучающие данные
            test_data: Тестовые данные

        Returns:
            Отчет об уязвимостях
        """
        if self.report is None:
            self.start_audit()

        # Вычисление ошибок
        train_errors = []
        for x, target in train_data:
            pred = self.core.signal(x)
            train_errors.append((pred - target) ** 2)

        test_errors = []
        for x, target in test_data:
            pred = self.core.signal(x)
            test_errors.append((pred - target) ** 2)

        train_mse = np.mean(train_errors)
        test_mse = np.mean(test_errors)

        # Отношение тестовой ошибки к обучающей
        if train_mse > 0:
            generalization_gap = test_mse / train_mse

            if generalization_gap > 5:
                self.report.add_vulnerability(
                    name='Overfitting',
                    severity='high',
                    description='Модель переобучена на обучающих данных',
                    details={
                        'train_mse': train_mse,
                        'test_mse': test_mse,
                        'generalization_gap': generalization_gap
                    }
                )
                self.report.add_recommendation(
                    'Добавьте регуляризацию, dropout или увеличьте размер данных'
                )

        return self.report

    def run_full_audit(self, train_data: Optional[List[Tuple[torch.Tensor, float]]] = None,
                       test_data: Optional[List[Tuple[torch.Tensor, float]]] = None
                       ) -> VulnerabilityReport:
        """
        Запуск полного аудита.

        Args:
            train_data: Обучающие данные
            test_data: Тестовые данные

        Returns:
            Полный отчет об уязвимостях
        """
        logger.info("Starting full vulnerability audit...")

        self.start_audit()
        self.check_weight_anomalies()
        self.check_gradient_sensitivity()
        self.check_adversarial_vulnerability()
        self.check_numerical_stability()

        if train_data and test_data:
            self.check_overfitting(train_data, test_data)

        logger.info(f"Audit completed. Risk score: {self.report.risk_score:.2f}")

        return self.report

    def get_security_rating(self) -> str:
        """
        Получение рейтинга безопасности.

        Returns:
            Рейтинг ('A', 'B', 'C', 'D', 'F')
        """
        if self.report is None:
            return 'N/A'

        if self.report.risk_score < 10:
            return 'A'
        elif self.report.risk_score < 25:
            return 'B'
        elif self.report.risk_score < 50:
            return 'C'
        elif self.report.risk_score < 75:
            return 'D'
        else:
            return 'F'


def audit_model(core: KokaoCore, 
                save_path: Optional[str] = None) -> VulnerabilityReport:
    """
    Быстрый аудит модели.

    Args:
        core: Модель для аудита
        save_path: Путь для сохранения отчета

    Returns:
        Отчет об уязвимостях
    """
    auditor = VulnerabilityAuditor(core)
    report = auditor.run_full_audit()

    if save_path:
        report.save(save_path)

    return report
