"""  
端口扫描 Flow

负责编排端口扫描的完整流程

架构：
- Flow 负责编排多个原子 Task
- 支持串行执行扫描工具（流式处理）
- 每个 Task 可独立重试
- 配置由 YAML 解析
"""

# Django 环境初始化（导入即生效）
from apps.common.prefect_django_setup import setup_django_for_prefect

import logging
import os
import subprocess
from pathlib import Path
from typing import Callable
from prefect import flow
from apps.scan.tasks.port_scan import (
    export_scan_targets_task,
    run_and_stream_save_ports_task
)
from apps.scan.handlers.scan_flow_handlers import (
    on_scan_flow_running,
    on_scan_flow_completed,
    on_scan_flow_failed,
)
from apps.scan.utils import config_parser, build_scan_command

logger = logging.getLogger(__name__)


def calculate_port_scan_timeout(
    tool_config: dict,
    file_path: str,
    base_per_pair: float = 0.5
) -> int:
    """
    根据目标数量和端口数量计算超时时间
    
    计算公式：超时时间 = 目标数 × 端口数 × base_per_pair
    超时范围：60秒 ~ 2天（172800秒）
    
    Args:
        tool_config: 工具配置字典，包含端口配置（ports, top-ports等）
        file_path: 目标文件路径（域名/IP列表）
        base_per_pair: 每个"端口-目标对"的基础时间（秒），默认 0.5秒
    
    Returns:
        int: 计算出的超时时间（秒），范围：60 ~ 172800
    
    Example:
        # 100个目标 × 100个端口 × 0.5秒 = 5000秒
        # 10个目标 × 1000个端口 × 0.5秒 = 5000秒
        timeout = calculate_port_scan_timeout(
            tool_config={'top-ports': 100},
            file_path='/path/to/domains.txt'
        )
    """
    try:
        # 1. 统计目标数量
        result = subprocess.run(
            ['wc', '-l', file_path],
            capture_output=True,
            text=True,
            check=True
        )
        target_count = int(result.stdout.strip().split()[0])
        
        # 2. 解析端口数量
        port_count = _parse_port_count(tool_config)
        
        # 3. 计算超时时间
        # 总工作量 = 目标数 × 端口数
        total_work = target_count * port_count
        timeout = int(total_work * base_per_pair)
        
        # 4. 设置合理的下限（不再设置上限）
        min_timeout = 60       # 最小 60 秒
        timeout = max(min_timeout, timeout)
        
        logger.info(
            f"计算端口扫描 timeout - "
            f"目标数: {target_count}, "
            f"端口数: {port_count}, "
            f"总工作量: {total_work}, "
            f"超时: {timeout}秒"
        )
        return timeout
        
    except Exception as e:
        logger.warning(f"计算 timeout 失败: {e}，使用默认值 600秒")
        return 600


def _parse_port_count(tool_config: dict) -> int:
    """
    从工具配置中解析端口数量
    
    优先级：
    1. top-ports: N  → 返回 N
    2. ports: "80,443,8080"  → 返回逗号分隔的数量
    3. ports: "1-1000"  → 返回范围的大小
    4. ports: "1-65535"  → 返回 65535
    5. 默认  → 返回 100（naabu 默认扫描 top 100）
    
    Args:
        tool_config: 工具配置字典
    
    Returns:
        int: 端口数量
    """
    # 1. 检查 top-ports 配置
    if 'top-ports' in tool_config:
        top_ports = tool_config['top-ports']
        if isinstance(top_ports, int) and top_ports > 0:
            return top_ports
        logger.warning(f"top-ports 配置无效: {top_ports}，使用默认值")
    
    # 2. 检查 ports 配置
    if 'ports' in tool_config:
        ports_str = str(tool_config['ports']).strip()
        
        # 2.1 逗号分隔的端口列表：80,443,8080
        if ',' in ports_str:
            port_list = [p.strip() for p in ports_str.split(',') if p.strip()]
            return len(port_list)
        
        # 2.2 端口范围：1-1000
        if '-' in ports_str:
            try:
                start, end = ports_str.split('-', 1)
                start_port = int(start.strip())
                end_port = int(end.strip())
                
                if 1 <= start_port <= end_port <= 65535:
                    return end_port - start_port + 1
                logger.warning(f"端口范围无效: {ports_str}，使用默认值")
            except ValueError:
                logger.warning(f"端口范围解析失败: {ports_str}，使用默认值")
        
        # 2.3 单个端口
        try:
            port = int(ports_str)
            if 1 <= port <= 65535:
                return 1
        except ValueError:
            logger.warning(f"端口配置解析失败: {ports_str}，使用默认值")
    
    # 3. 默认值：naabu 默认扫描 top 100 端口
    return 100





def _export_scan_targets(target_id: int, port_scan_dir: Path) -> tuple[str, int, str]:
    """
    导出扫描目标到文件
    
    根据 Target 类型自动决定导出内容：
    - DOMAIN: 从 Subdomain 表导出子域名
    - IP: 直接写入 target.name
    - CIDR: 展开 CIDR 范围内的所有 IP
    
    Args:
        target_id: 目标 ID
        port_scan_dir: 端口扫描目录
        
    Returns:
        tuple: (targets_file, target_count, target_type)
    """
    logger.info("Step 1: 导出扫描目标列表")
    
    targets_file = str(port_scan_dir / 'targets.txt')
    export_result = export_scan_targets_task(
        target_id=target_id,
        output_file=targets_file,
        batch_size=1000  # 每次读取 1000 条，优化内存占用
    )
    
    target_count = export_result['total_count']
    target_type = export_result.get('target_type', 'unknown')
    
    logger.info(
        "✓ 扫描目标导出完成 - 类型: %s, 文件: %s, 数量: %d",
        target_type,
        export_result['output_file'],
        target_count
    )
    
    if target_count == 0:
        logger.warning("目标下没有可扫描的地址，无法执行端口扫描")
    
    return export_result['output_file'], target_count, target_type


def _run_scans_sequentially(
    enabled_tools: dict,
    domains_file: str,
    port_scan_dir: Path,
    scan_id: int,
    target_id: int,
    target_name: str
) -> tuple[dict, int, list, list]:
    """
    串行执行端口扫描任务
    
    Args:
        enabled_tools: 已启用的工具配置字典
        domains_file: 域名文件路径
        port_scan_dir: 端口扫描目录
        scan_id: 扫描任务 ID
        target_id: 目标 ID
        target_name: 目标名称（用于错误日志）
        
    Returns:
        tuple: (tool_stats, processed_records, successful_tool_names, failed_tools)
        注意：端口扫描是流式输出，不生成结果文件
        
    Raises:
        RuntimeError: 所有工具均失败
    """
    # ==================== 构建命令并串行执行 ====================
    
    tool_stats = {}
    processed_records = 0
    failed_tools = []      # 记录失败的工具（含原因）
    
    # for循环执行工具：按顺序串行运行每个启用的端口扫描工具
    for tool_name, tool_config in enabled_tools.items():
        # 1. 构建完整命令（变量替换）
        try:
            command = build_scan_command(
                tool_name=tool_name,
                scan_type='port_scan',
                command_params={
                    'domains_file': domains_file  # 对应 {domains_file}
                },
                tool_config=tool_config     #yaml的工具配置
            )
        except Exception as e:
            reason = f"命令构建失败: {str(e)}"
            logger.error(f"构建 {tool_name} 命令失败: {e}")
            failed_tools.append({'tool': tool_name, 'reason': reason})
            continue
        
        # 2. 获取超时时间（支持 'auto' 动态计算）
        config_timeout = tool_config['timeout']
        if config_timeout == 'auto':
            # 动态计算超时时间
            config_timeout = calculate_port_scan_timeout(
                tool_config=tool_config,
                file_path=str(domains_file)
            )
            logger.info(f"✓ 工具 {tool_name} 动态计算 timeout: {config_timeout}秒")
        
        # 2.1 生成日志文件路径
        from datetime import datetime
        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
        log_file = port_scan_dir / f"{tool_name}_{timestamp}.log"
        
        # 3. 执行扫描任务
        logger.info("开始执行 %s 扫描（超时: %d秒）...", tool_name, config_timeout)
        
        try:
            # 直接调用 task（串行执行）
            # 注意：端口扫描是流式输出到 stdout，不使用 output_file
            result = run_and_stream_save_ports_task(
                cmd=command,
                tool_name=tool_name,  # 工具名称
                scan_id=scan_id,
                target_id=target_id,
                cwd=str(port_scan_dir),
                shell=True,
                batch_size=1000,
                timeout=config_timeout,
                log_file=str(log_file)  # 新增：日志文件路径
            )
            
            tool_stats[tool_name] = {
                'command': command,
                'result': result,
                'timeout': config_timeout
            }
            processed_records += result.get('processed_records', 0)
            logger.info(
                "✓ 工具 %s 流式处理完成 - 记录数: %d",
                tool_name, result.get('processed_records', 0)
            )
            
        except subprocess.TimeoutExpired as exc:
            # 超时异常单独处理
            # 注意：流式处理任务超时时，已解析的数据已保存到数据库
            reason = f"执行超时（配置: {config_timeout}秒）"
            failed_tools.append({'tool': tool_name, 'reason': reason})
            logger.warning(
                "⚠️ 工具 %s 执行超时 - 超时配置: %d秒\n"
                "注意：超时前已解析的端口数据已保存到数据库，但扫描未完全完成。",
                tool_name, config_timeout
            )
        except Exception as exc:
            # 其他异常
            failed_tools.append({'tool': tool_name, 'reason': str(exc)})
            logger.error("工具 %s 执行失败: %s", tool_name, exc, exc_info=True)
    
    if failed_tools:
        logger.warning(
            "以下扫描工具执行失败: %s",
            ', '.join([f['tool'] for f in failed_tools])
        )
    
    if not tool_stats:
        error_details = "; ".join([f"{f['tool']}: {f['reason']}" for f in failed_tools])
        logger.warning("所有端口扫描工具均失败 - 目标: %s, 失败工具: %s", target_name, error_details)
        # 返回空结果，不抛出异常，让扫描继续
        return {}, 0, [], failed_tools
    
    # 动态计算成功的工具列表
    successful_tool_names = [name for name in enabled_tools.keys() 
                              if name not in [f['tool'] for f in failed_tools]]
    
    logger.info(
        "✓ 串行端口扫描执行完成 - 成功: %d/%d (成功: %s, 失败: %s)",
        len(tool_stats), len(enabled_tools),
        ', '.join(successful_tool_names) if successful_tool_names else '无',
        ', '.join([f['tool'] for f in failed_tools]) if failed_tools else '无'
    )
    
    return tool_stats, processed_records, successful_tool_names, failed_tools


@flow(
    name="port_scan", 
    log_prints=True,
    on_running=[on_scan_flow_running],
    on_completion=[on_scan_flow_completed],
    on_failure=[on_scan_flow_failed],
)
def port_scan_flow(
    scan_id: int,
    target_name: str,
    target_id: int,
    scan_workspace_dir: str,
    enabled_tools: dict
) -> dict:
    """
    端口扫描 Flow
    
    主要功能：
        1. 扫描目标域名/IP 的开放端口
        2. 保存 host + ip + port 三元映射到 HostPortMapping 表
    
    输出资产：
        - HostPortMapping：主机端口映射（host + ip + port 三元组）
    
    工作流程：
        Step 0: 创建工作目录
        Step 1: 导出域名列表到文件（供扫描工具使用）
        Step 2: 解析配置，获取启用的工具
        Step 3: 串行执行扫描工具，运行端口扫描工具并实时解析输出到数据库（→ HostPortMapping）

    Args:
        scan_id: 扫描任务 ID
        target_name: 域名
        target_id: 目标 ID
        scan_workspace_dir: Scan 工作空间目录
        enabled_tools: 启用的工具配置字典

    Returns:
        dict: {
            'success': bool,
            'scan_id': int,
            'target': str,
            'scan_workspace_dir': str,
            'domains_file': str,
            'domain_count': int,
            'processed_records': int,
            'executed_tasks': list,
            'tool_stats': {
                'total': int,                    # 总工具数
                'successful': int,               # 成功工具数
                'failed': int,                   # 失败工具数
                'successful_tools': list[str],   # 成功工具列表 ['naabu_active']
                'failed_tools': list[dict],      # 失败工具列表 [{'tool': 'naabu_passive', 'reason': '超时'}]
                'details': dict                  # 详细执行结果（保留向后兼容）
            }
        }

    Raises:
        ValueError: 配置错误
        RuntimeError: 执行失败
    
    Note:
        端口扫描工具（如 naabu）会解析域名获取 IP，输出 host + ip + port 三元组。
        同一 host 可能对应多个 IP（CDN、负载均衡），因此使用三元映射表存储。
    """
    try:
        # 参数验证
        if scan_id is None:
            raise ValueError("scan_id 不能为空")
        if not target_name:
            raise ValueError("target_name 不能为空")
        if target_id is None:
            raise ValueError("target_id 不能为空")
        if not scan_workspace_dir:
            raise ValueError("scan_workspace_dir 不能为空")
        if not enabled_tools:
            raise ValueError("enabled_tools 不能为空")
        
        logger.info(
            "="*60 + "\n" +
            "开始端口扫描\n" +
            f"  Scan ID: {scan_id}\n" +
            f"  Target: {target_name}\n" +
            f"  Workspace: {scan_workspace_dir}\n" +
            "="*60
        )
        
        # Step 0: 创建工作目录
        from apps.scan.utils import setup_scan_directory
        port_scan_dir = setup_scan_directory(scan_workspace_dir, 'port_scan')
        
        # Step 1: 导出扫描目标列表到文件（根据 Target 类型自动决定内容）
        targets_file, target_count, target_type = _export_scan_targets(target_id, port_scan_dir)
        
        if target_count == 0:
            logger.warning("目标下没有可扫描的地址，跳过端口扫描")
            return {
                'success': True,
                'scan_id': scan_id,
                'target': target_name,
                'scan_workspace_dir': scan_workspace_dir,
                'targets_file': targets_file,
                'target_count': 0,
                'target_type': target_type,
                'processed_records': 0,
                'executed_tasks': ['export_scan_targets'],
                'tool_stats': {
                    'total': 0,
                    'successful': 0,
                    'failed': 0,
                    'successful_tools': [],
                    'failed_tools': [],
                    'details': {}
                }
            }
        
        # Step 2: 工具配置信息
        logger.info("Step 2: 工具配置信息")
        logger.info(
            "✓ 启用工具: %s",
            ', '.join(enabled_tools.keys())
        )
        
        # Step 3: 串行执行扫描工具
        logger.info("Step 3: 串行执行扫描工具")
        tool_stats, processed_records, successful_tool_names, failed_tools = _run_scans_sequentially(
            enabled_tools=enabled_tools,
            domains_file=targets_file,  # 现在是 targets_file，兼容原参数名
            port_scan_dir=port_scan_dir,
            scan_id=scan_id,
            target_id=target_id,
            target_name=target_name
        )
        
        logger.info("="*60 + "\n✓ 端口扫描完成\n" + "="*60)
        
        # 动态生成已执行的任务列表
        executed_tasks = ['export_scan_targets', 'parse_config']
        executed_tasks.extend([f'run_and_stream_save_ports ({tool})' for tool in tool_stats.keys()])
        
        return {
            'success': True,
            'scan_id': scan_id,
            'target': target_name,
            'scan_workspace_dir': scan_workspace_dir,
            'targets_file': targets_file,
            'target_count': target_count,
            'target_type': target_type,
            'processed_records': processed_records,
            'executed_tasks': executed_tasks,
            'tool_stats': {
                'total': len(tool_stats) + len(failed_tools),
                'successful': len(successful_tool_names),
                'failed': len(failed_tools),
                'successful_tools': successful_tool_names,
                'failed_tools': failed_tools,  # [{'tool': 'naabu_active', 'reason': '超时'}]
                'details': tool_stats  # 详细结果（保留向后兼容）
            }
        }

    except ValueError as e:
        logger.error("配置错误: %s", e)
        raise
    except RuntimeError as e:
        logger.error("运行时错误: %s", e)
        raise
    except Exception as e:
        logger.exception("端口扫描失败: %s", e)
        raise
