# Copyright (c) 2024-2026 Tencent Zhuque Lab. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Requirement: Any integration or derivative work must explicitly attribute
# Tencent Zhuque Lab (https://github.com/Tencent/AI-Infra-Guard) in its
# documentation or user interface, as detailed in the NOTICE file.

from typing import Dict, List, Optional, Any, Union
from pathlib import Path
import json
from loguru import logger

from .plugin_validator import PluginValidator
from .plugin_registry import PluginRegistry


class PluginLoader:
    """插件加载器，负责扫描和加载插件"""
    
    def __init__(self, registry: Optional[PluginRegistry] = None):
        self.validator = PluginValidator()
        self.registry = registry or PluginRegistry()
        self.default_plugin_dirs = [
            Path("./plugin"),
            Path("./custom_plugins"),
            Path.home() / ".deepteam" / "plugins"
        ]
    
    def load_plugin(self, plugin_path: Union[str, Path]) -> Dict[str, Any]:
        """加载单个插件"""
        plugin_path = Path(plugin_path)
        
        result = {
            'success': False,
            'plugin_info': None,
            'errors': [],
            'warnings': []
        }
        
        try:
            # 验证插件
            if plugin_path.is_file():
                validation_result = self.validator.validate_plugin_file(plugin_path)
            elif plugin_path.is_dir():
                validation_result = self.validator.validate_plugin_directory(plugin_path)
            else:
                result['errors'].append(f"路径不存在: {plugin_path}")
                return result
            
            # 处理验证结果
            if validation_result['valid']:
                if plugin_path.is_file():
                    # 单个文件
                    if self.registry.register_plugin(plugin_path, validation_result):
                        result['success'] = True
                        result['plugin_info'] = validation_result
                    else:
                        result['errors'].append("插件注册失败")
                else:
                    # 目录
                    success_count = 0
                    for plugin_info in validation_result['plugins']:
                        # 为每个插件找到对应的文件
                        plugin_file = self._find_plugin_file(plugin_path, plugin_info['class_name'])
                        if plugin_file and self.registry.register_plugin(plugin_file, plugin_info):
                            success_count += 1
                    
                    if success_count > 0:
                        result['success'] = True
                        result['plugin_info'] = {
                            'type': 'directory',
                            'loaded_count': success_count,
                            'total_count': len(validation_result['plugins'])
                        }
                    else:
                        result['errors'].append("没有成功加载任何插件")
            else:
                result['errors'].extend(validation_result['errors'])
                result['warnings'].extend(validation_result['warnings'])
                
        except Exception as e:
            result['errors'].append(f"加载插件时发生错误: {str(e)}")
        
        return result
    
    def _find_plugin_file(self, dir_path: Path, class_name: str) -> Optional[Path]:
        """在目录中查找包含指定类的文件"""
        for py_file in dir_path.glob("*.py"):
            try:
                validation_result = self.validator.validate_plugin_file(py_file)
                if validation_result['valid'] and validation_result['class_name'] == class_name:
                    return py_file
            except Exception:
                continue
        return None
    
    def load_plugins_from_directory(self, dir_path: Union[str, Path]) -> Dict[str, Any]:
        """从目录加载所有插件"""
        dir_path = Path(dir_path)
        
        result = {
            'success': False,
            'loaded_plugins': [],
            'errors': [],
            'warnings': []
        }
        
        try:
            if not dir_path.exists() or not dir_path.is_dir():
                result['errors'].append(f"目录不存在或不是有效目录: {dir_path}")
                return result
            
            # 验证目录
            validation_result = self.validator.validate_plugin_directory(dir_path)
            
            if validation_result['valid']:
                success_count = 0
                for plugin_info in validation_result['plugins']:
                    plugin_file = self._find_plugin_file(dir_path, plugin_info['class_name'])
                    if plugin_file and self.registry.register_plugin(plugin_file, plugin_info):
                        success_count += 1
                        result['loaded_plugins'].append({
                            'file': str(plugin_file),
                            'class': plugin_info['class_name'],
                            'type': plugin_info['plugin_type']
                        })
                
                if success_count > 0:
                    result['success'] = True
                    logger.info(f"从目录 {dir_path} 成功加载了 {success_count} 个插件")
                else:
                    result['errors'].append("没有成功加载任何插件")
            else:
                result['errors'].extend(validation_result['errors'])
                result['warnings'].extend(validation_result['warnings'])
                
        except Exception as e:
            result['errors'].append(f"加载插件目录时发生错误: {str(e)}")
        
        return result
    
    def auto_discover_plugins(self) -> Dict[str, Any]:
        """自动发现并加载插件"""
        result = {
            'success': False,
            'discovered_dirs': [],
            'loaded_plugins': [],
            'errors': [],
            'warnings': []
        }
        
        for plugin_dir in self.default_plugin_dirs:
            if plugin_dir.exists() and plugin_dir.is_dir():
                result['discovered_dirs'].append(str(plugin_dir))
                dir_result = self.load_plugins_from_directory(plugin_dir)
                
                if dir_result['success']:
                    result['loaded_plugins'].extend(dir_result['loaded_plugins'])
                else:
                    result['errors'].extend(dir_result['errors'])
                    result['warnings'].extend(dir_result['warnings'])
        
        result['success'] = len(result['loaded_plugins']) > 0
        
        if result['success']:
            logger.info(f"自动发现并加载了 {len(result['loaded_plugins'])} 个插件")
        
        return result
    
    def load_plugins_from_config(self, config_file: Union[str, Path]) -> Dict[str, Any]:
        """从配置文件加载插件"""
        config_file = Path(config_file)
        
        result = {
            'success': False,
            'loaded_plugins': [],
            'errors': [],
            'warnings': []
        }
        
        try:
            if not config_file.exists():
                result['errors'].append(f"配置文件不存在: {config_file}")
                return result
            
            with open(config_file, 'r', encoding='utf-8') as f:
                config = json.load(f)
            
            # 处理插件配置
            plugins_config = config.get('plugins', [])
            
            for plugin_config in plugins_config:
                plugin_path = plugin_config.get('path')
                if not plugin_path:
                    result['errors'].append("插件配置缺少path字段")
                    continue
                
                plugin_result = self.load_plugin(plugin_path)
                if plugin_result['success']:
                    result['loaded_plugins'].append({
                        'path': plugin_path,
                        'info': plugin_result['plugin_info']
                    })
                else:
                    result['errors'].extend(plugin_result['errors'])
                    result['warnings'].extend(plugin_result['warnings'])
            
            result['success'] = len(result['loaded_plugins']) > 0
            
        except json.JSONDecodeError as e:
            result['errors'].append(f"配置文件格式错误: {e}")
        except Exception as e:
            result['errors'].append(f"加载配置文件时发生错误: {str(e)}")
        
        return result
    
    def get_loaded_plugins(self) -> Dict[str, List[str]]:
        """获取已加载的插件列表"""
        return self.registry.list_plugins()
    
    def get_plugin_info(self, plugin_key: str) -> Optional[Dict[str, Any]]:
        """获取插件详细信息"""
        return self.registry.get_plugin_info(plugin_key)
    
    def unload_plugin(self, plugin_key: str) -> bool:
        """卸载插件"""
        return self.registry.unregister_plugin(plugin_key)
    
    def reload_plugin(self, plugin_path: Union[str, Path]) -> Dict[str, Any]:
        """重新加载插件"""
        plugin_path = Path(plugin_path)
        
        # 先卸载已存在的插件
        plugin_info = self.registry.get_plugin_info(f"attack_{plugin_path.stem}")
        if not plugin_info:
            plugin_info = self.registry.get_plugin_info(f"metric_{plugin_path.stem}")
        
        if plugin_info:
            self.registry.unregister_plugin(f"{plugin_info['type']}_{plugin_info['class_name']}")
        
        # 重新加载
        return self.load_plugin(plugin_path) 