"""
目录扫描 Flow

负责编排目录扫描的完整流程

架构：
- Flow 负责编排多个原子 Task
- 支持并发执行扫描工具（使用 ThreadPoolTaskRunner）
- 每个 Task 可独立重试
- 配置由 YAML 解析
"""

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

from prefect import flow
from prefect.task_runners import ThreadPoolTaskRunner

import hashlib
import logging
import os
import subprocess
from datetime import datetime
from pathlib import Path
from typing import List, Tuple

from apps.scan.tasks.directory_scan import (
    export_sites_task,
    run_and_stream_save_directories_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, ensure_wordlist_local

logger = logging.getLogger(__name__)

# 默认最大并发数
DEFAULT_MAX_WORKERS = 5


def calculate_directory_scan_timeout(
    tool_config: dict,
    base_per_word: float = 1.0,
    min_timeout: int = 60,
    max_timeout: int = 7200
) -> int:
    """
    根据字典行数计算目录扫描超时时间
    
    计算公式：超时时间 = 字典行数 × 每个单词基础时间
    超时范围：60秒 ~ 2小时（7200秒）
    
    Args:
        tool_config: 工具配置字典，包含 wordlist 路径
        base_per_word: 每个单词的基础时间（秒），默认 1.0秒
        min_timeout: 最小超时时间（秒），默认 60秒
        max_timeout: 最大超时时间（秒），默认 7200秒（2小时）
    
    Returns:
        int: 计算出的超时时间（秒），范围：60 ~ 7200
    
    Example:
        # 1000行字典 × 1.0秒 = 1000秒 → 限制为7200秒中的 1000秒
        # 10000行字典 × 1.0秒 = 10000秒 → 限制为7200秒（最大值）
        timeout = calculate_directory_scan_timeout(
            tool_config={'wordlist': '/path/to/wordlist.txt'}
        )
    """
    try:
        # 从 tool_config 中获取 wordlist 路径
        wordlist_path = tool_config.get('wordlist')
        if not wordlist_path:
            logger.warning("工具配置中未指定 wordlist，使用默认超时: %d秒", min_timeout)
            return min_timeout
        
        # 展开用户目录（~）
        wordlist_path = os.path.expanduser(wordlist_path)
        
        # 检查文件是否存在
        if not os.path.exists(wordlist_path):
            logger.warning("字典文件不存在: %s，使用默认超时: %d秒", wordlist_path, min_timeout)
            return min_timeout
        
        # 使用 wc -l 快速统计字典行数
        result = subprocess.run(
            ['wc', '-l', wordlist_path],
            capture_output=True,
            text=True,
            check=True
        )
        # wc -l 输出格式：行数 + 空格 + 文件名
        line_count = int(result.stdout.strip().split()[0])
        
        # 计算超时时间
        timeout = int(line_count * base_per_word)
        
        # 设置合理的下限（不再设置上限）
        timeout = max(min_timeout, timeout)
        
        logger.info(
            "目录扫描超时计算 - 字典: %s, 行数: %d, 基础时间: %.3f秒/词, 计算超时: %d秒",
            wordlist_path, line_count, base_per_word, timeout
        )
        
        return timeout
        
    except subprocess.CalledProcessError as e:
        logger.error("统计字典行数失败: %s", e)
        # 失败时返回默认超时
        return min_timeout
    except (ValueError, IndexError) as e:
        logger.error("解析字典行数失败: %s", e)
        return min_timeout
    except Exception as e:
        logger.error("计算超时时间异常: %s", e)
        return min_timeout


def _get_max_workers(tool_config: dict, default: int = DEFAULT_MAX_WORKERS) -> int:
    """
    从单个工具配置中获取 max_workers 参数
    
    Args:
        tool_config: 单个工具的配置字典，如 {'max_workers': 10, 'threads': 5, ...}
        default: 默认值，默认为 5
        
    Returns:
        int: max_workers 值
    """
    if not isinstance(tool_config, dict):
        return default
    
    # 支持 max_workers 和 max-workers（YAML 中划线会被转换）
    max_workers = tool_config.get('max_workers') or tool_config.get('max-workers')
    if max_workers is not None and isinstance(max_workers, int) and max_workers > 0:
        return max_workers
    return default





def _export_site_urls(target_id: int, target_name: str, directory_scan_dir: Path) -> tuple[str, int]:
    """
    导出目标下的所有站点 URL 到文件（支持懒加载）
    
    Args:
        target_id: 目标 ID
        target_name: 目标名称（用于懒加载创建默认站点）
        directory_scan_dir: 目录扫描目录
        
    Returns:
        tuple: (sites_file, site_count)
        
    Raises:
        ValueError: 站点数量为 0
    """
    logger.info("Step 1: 导出目标的所有站点 URL")
    
    sites_file = str(directory_scan_dir / 'sites.txt')
    export_result = export_sites_task(
        target_id=target_id,
        output_file=sites_file,
        batch_size=1000  # 每次读取 1000 条，优化内存占用
    )
    
    site_count = export_result['total_count']
    
    logger.info(
        "✓ 站点 URL 导出完成 - 文件: %s, 数量: %d",
        export_result['output_file'],
        site_count
    )
    
    if site_count == 0:
        logger.warning("目标下没有站点，无法执行目录扫描")
        # 不抛出异常，由上层决定如何处理
        # raise ValueError("目标下没有站点，无法执行目录扫描")
    
    return export_result['output_file'], site_count


def _run_scans_sequentially(
    enabled_tools: dict,
    sites_file: str,
    directory_scan_dir: Path,
    scan_id: int,
    target_id: int,
    site_count: int,
    target_name: str
) -> tuple[int, int, list]:
    """
    串行执行目录扫描任务（支持多工具）- 已废弃，保留用于兼容
    
    Args:
        enabled_tools: 启用的工具配置字典
        sites_file: 站点文件路径
        directory_scan_dir: 目录扫描目录
        scan_id: 扫描任务 ID
        target_id: 目标 ID
        site_count: 站点数量
        target_name: 目标名称（用于错误日志）
        
    Returns:
        tuple: (total_directories, processed_sites, failed_sites)
    """
    # 读取站点列表
    sites = []
    with open(sites_file, 'r', encoding='utf-8') as f:
        for line in f:
            site_url = line.strip()
            if site_url:
                sites.append(site_url)
    
    logger.info("准备扫描 %d 个站点，使用工具: %s", len(sites), ', '.join(enabled_tools.keys()))
    
    total_directories = 0
    processed_sites_set = set()  # 使用 set 避免重复计数
    failed_sites = []
    
    # 遍历每个工具
    for tool_name, tool_config in enabled_tools.items():
        logger.info("="*60)
        logger.info("使用工具: %s", tool_name)
        logger.info("="*60)

        # 如果配置了 wordlist_name，则先确保本地存在对应的字典文件（含 hash 校验）
        wordlist_name = tool_config.get('wordlist_name')
        if wordlist_name:
            try:
                local_wordlist_path = ensure_wordlist_local(wordlist_name)
                tool_config['wordlist'] = local_wordlist_path
            except Exception as exc:
                logger.error("为工具 %s 准备字典失败: %s", tool_name, exc)
                # 当前工具无法执行，将所有站点视为失败，继续下一个工具
                failed_sites.extend(sites)
                continue
        
        # 逐个站点执行扫描
        for idx, site_url in enumerate(sites, 1):
            logger.info(
                "[%d/%d] 开始扫描站点: %s (工具: %s)",
                idx, len(sites), site_url, tool_name
            )
            
            # 使用统一的命令构建器
            try:
                command = build_scan_command(
                    tool_name=tool_name,
                    scan_type='directory_scan',
                    command_params={
                        'url': site_url
                    },
                    tool_config=tool_config
                )
            except Exception as e:
                logger.error(
                    "✗ [%d/%d] 构建 %s 命令失败: %s - 站点: %s",
                    idx, len(sites), tool_name, e, site_url
                )
                failed_sites.append(site_url)
                continue
            
            # 单个站点超时：从配置中获取（支持 'auto' 动态计算）
            # ffuf 逐个站点扫描，timeout 就是单个站点的超时时间
            site_timeout = tool_config.get('timeout', 300)
            if site_timeout == 'auto':
                # 动态计算超时时间（基于字典行数）
                site_timeout = calculate_directory_scan_timeout(tool_config)
                logger.info(f"✓ 工具 {tool_name} 动态计算 timeout: {site_timeout}秒")
            
            # 生成日志文件路径
            from datetime import datetime
            timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
            log_file = directory_scan_dir / f"{tool_name}_{timestamp}_{idx}.log"
            
            try:
                # 直接调用 task（串行执行）
                result = run_and_stream_save_directories_task(
                    cmd=command,
                    tool_name=tool_name,  # 新增：工具名称
                    scan_id=scan_id,
                    target_id=target_id,
                    site_url=site_url,
                    cwd=str(directory_scan_dir),
                    shell=True,
                    batch_size=1000,
                    timeout=site_timeout,
                    log_file=str(log_file)  # 新增：日志文件路径
                )
                
                total_directories += result.get('created_directories', 0)
                processed_sites_set.add(site_url)  # 使用 set 记录成功的站点
                
                logger.info(
                    "✓ [%d/%d] 站点扫描完成: %s - 发现 %d 个目录",
                    idx, len(sites), site_url,
                    result.get('created_directories', 0)
                )
                
            except subprocess.TimeoutExpired as exc:
                # 超时异常单独处理
                failed_sites.append(site_url)
                logger.warning(
                    "⚠️ [%d/%d] 站点扫描超时: %s - 超时配置: %d秒\n"
                    "注意：超时前已解析的目录数据已保存到数据库，但扫描未完全完成。",
                    idx, len(sites), site_url, site_timeout
                )
            except Exception as exc:
                # 其他异常
                failed_sites.append(site_url)
                logger.error(
                    "✗ [%d/%d] 站点扫描失败: %s - 错误: %s",
                    idx, len(sites), site_url, exc
                )
            
            # 每 10 个站点输出进度
            if idx % 10 == 0:
                logger.info(
                    "进度: %d/%d (%.1f%%) - 已发现 %d 个目录",
                    idx, len(sites), idx/len(sites)*100, total_directories
                )
    
    # 计算成功和失败的站点数
    processed_count = len(processed_sites_set)
    
    if failed_sites:
        logger.warning(
            "部分站点扫描失败: %d/%d",
            len(failed_sites), len(sites)
        )
    
    logger.info(
        "✓ 串行目录扫描执行完成 - 成功: %d/%d, 失败: %d, 总目录数: %d",
        processed_count, len(sites), len(failed_sites), total_directories
    )
    
    return total_directories, processed_count, failed_sites


def _generate_log_filename(tool_name: str, site_url: str, directory_scan_dir: Path) -> Path:
    """
    生成唯一的日志文件名
    
    使用 URL 的 hash 确保并发时不会冲突
    
    Args:
        tool_name: 工具名称
        site_url: 站点 URL
        directory_scan_dir: 目录扫描目录
        
    Returns:
        Path: 日志文件路径
    """
    url_hash = hashlib.md5(site_url.encode()).hexdigest()[:8]
    timestamp = datetime.now().strftime('%Y%m%d_%H%M%S_%f')
    return directory_scan_dir / f"{tool_name}_{url_hash}_{timestamp}.log"


def _run_scans_concurrently(
    enabled_tools: dict,
    sites_file: str,
    directory_scan_dir: Path,
    scan_id: int,
    target_id: int,
    site_count: int,
    target_name: str
) -> Tuple[int, int, List[str]]:
    """
    并发执行目录扫描任务（使用 ThreadPoolTaskRunner）
    
    Args:
        enabled_tools: 启用的工具配置字典
        sites_file: 站点文件路径
        directory_scan_dir: 目录扫描目录
        scan_id: 扫描任务 ID
        target_id: 目标 ID
        site_count: 站点数量
        target_name: 目标名称（用于错误日志）
        
    Returns:
        tuple: (total_directories, processed_sites, failed_sites)
    """
    # 读取站点列表
    sites: List[str] = []
    with open(sites_file, 'r', encoding='utf-8') as f:
        for line in f:
            site_url = line.strip()
            if site_url:
                sites.append(site_url)
    
    if not sites:
        logger.warning("站点列表为空")
        return 0, 0, []
    
    logger.info(
        "准备并发扫描 %d 个站点，使用工具: %s",
        len(sites), ', '.join(enabled_tools.keys())
    )
    
    total_directories = 0
    processed_sites_count = 0
    failed_sites: List[str] = []
    
    # 遍历每个工具
    for tool_name, tool_config in enabled_tools.items():
        # 每个工具独立获取 max_workers 配置
        max_workers = _get_max_workers(tool_config)
        
        logger.info("="*60)
        logger.info("使用工具: %s (并发模式, max_workers=%d)", tool_name, max_workers)
        logger.info("="*60)

        # 如果配置了 wordlist_name，则先确保本地存在对应的字典文件（含 hash 校验）
        wordlist_name = tool_config.get('wordlist_name')
        if wordlist_name:
            try:
                local_wordlist_path = ensure_wordlist_local(wordlist_name)
                tool_config['wordlist'] = local_wordlist_path
            except Exception as exc:
                logger.error("为工具 %s 准备字典失败: %s", tool_name, exc)
                # 当前工具无法执行，将所有站点视为失败，继续下一个工具
                failed_sites.extend(sites)
                continue
        
        # 计算超时时间（所有站点共用）
        site_timeout = tool_config.get('timeout', 300)
        if site_timeout == 'auto':
            site_timeout = calculate_directory_scan_timeout(tool_config)
            logger.info(f"✓ 工具 {tool_name} 动态计算 timeout: {site_timeout}秒")
        
        # 准备所有站点的扫描参数
        scan_params_list = []
        for idx, site_url in enumerate(sites, 1):
            try:
                command = build_scan_command(
                    tool_name=tool_name,
                    scan_type='directory_scan',
                    command_params={'url': site_url},
                    tool_config=tool_config
                )
                log_file = _generate_log_filename(tool_name, site_url, directory_scan_dir)
                scan_params_list.append({
                    'idx': idx,
                    'site_url': site_url,
                    'command': command,
                    'log_file': str(log_file),
                    'timeout': site_timeout
                })
            except Exception as e:
                logger.error(
                    "✗ [%d/%d] 构建 %s 命令失败: %s - 站点: %s",
                    idx, len(sites), tool_name, e, site_url
                )
                failed_sites.append(site_url)
        
        if not scan_params_list:
            logger.warning("没有有效的扫描任务")
            continue
        
        # ============================================================
        # 分批执行策略：控制实际并发的 ffuf 进程数
        # ============================================================
        total_tasks = len(scan_params_list)
        logger.info("开始分批执行 %d 个扫描任务（每批 %d 个）...", total_tasks, max_workers)
        
        batch_num = 0
        for batch_start in range(0, total_tasks, max_workers):
            batch_end = min(batch_start + max_workers, total_tasks)
            batch_params = scan_params_list[batch_start:batch_end]
            batch_num += 1
            
            logger.info("执行第 %d 批任务（%d-%d/%d）...", batch_num, batch_start + 1, batch_end, total_tasks)
            
            # 提交当前批次的任务（非阻塞，立即返回 future）
            futures = []
            for params in batch_params:
                future = run_and_stream_save_directories_task.submit(
                    cmd=params['command'],
                    tool_name=tool_name,
                    scan_id=scan_id,
                    target_id=target_id,
                    site_url=params['site_url'],
                    cwd=str(directory_scan_dir),
                    shell=True,
                    batch_size=1000,
                    timeout=params['timeout'],
                    log_file=params['log_file']
                )
                futures.append((params['idx'], params['site_url'], future))
            
            # 等待当前批次所有任务完成（阻塞，确保本批完成后再启动下一批）
            for idx, site_url, future in futures:
                try:
                    result = future.result()  # 阻塞等待单个任务完成
                    directories_found = result.get('created_directories', 0)
                    total_directories += directories_found
                    processed_sites_count += 1
                    
                    logger.info(
                        "✓ [%d/%d] 站点扫描完成: %s - 发现 %d 个目录",
                        idx, len(sites), site_url, directories_found
                    )
                    
                except Exception as exc:
                    failed_sites.append(site_url)
                    if 'timeout' in str(exc).lower() or isinstance(exc, subprocess.TimeoutExpired):
                        logger.warning(
                            "⚠️ [%d/%d] 站点扫描超时: %s - 错误: %s",
                            idx, len(sites), site_url, exc
                        )
                    else:
                        logger.error(
                            "✗ [%d/%d] 站点扫描失败: %s - 错误: %s",
                            idx, len(sites), site_url, exc
                        )
    
    # 输出汇总信息
    if failed_sites:
        logger.warning(
            "部分站点扫描失败: %d/%d",
            len(failed_sites), len(sites)
        )
    
    logger.info(
        "✓ 并发目录扫描执行完成 - 成功: %d/%d, 失败: %d, 总目录数: %d",
        processed_sites_count, len(sites), len(failed_sites), total_directories
    )
    
    return total_directories, processed_sites_count, failed_sites


@flow(
    name="directory_scan", 
    log_prints=True,
    on_running=[on_scan_flow_running],
    on_completion=[on_scan_flow_completed],
    on_failure=[on_scan_flow_failed],
)
def directory_scan_flow(
    scan_id: int,
    target_name: str,
    target_id: int,
    scan_workspace_dir: str,
    enabled_tools: dict
) -> dict:
    """
    目录扫描 Flow
    
    主要功能：
        1. 从 target 获取所有站点的 URL
        2. 对每个站点 URL 执行目录扫描（支持 ffuf 等工具）
        3. 流式保存扫描结果到数据库 Directory 表
    
    工作流程：
        Step 0: 创建工作目录
        Step 1: 导出站点 URL 列表到文件（供扫描工具使用）
        Step 2: 验证工具配置
        Step 3: 并发执行扫描工具并实时保存结果（使用 ThreadPoolTaskRunner）
    
    ffuf 输出字段：
        - url: 发现的目录/文件 URL
        - length: 响应内容长度
        - status: HTTP 状态码
        - words: 响应内容单词数
        - lines: 响应内容行数
        - content_type: 内容类型
        - duration: 请求耗时
    
    Args:
        scan_id: 扫描任务 ID
        target_name: 目标名称
        target_id: 目标 ID
        scan_workspace_dir: 扫描工作空间目录
        enabled_tools: 启用的工具配置字典
        
    Returns:
        dict: {
            'success': bool,
            'scan_id': int,
            'target': str,
            'scan_workspace_dir': str,
            'sites_file': str,
            'site_count': int,
            'total_directories': int,  # 发现的总目录数
            'processed_sites': int,  # 成功处理的站点数
            'failed_sites_count': int,  # 失败的站点数
            'executed_tasks': list
        }
    
    Raises:
        ValueError: 参数错误
        RuntimeError: 执行失败
    """
    try:
        logger.info(
            "="*60 + "\n" +
            "开始目录扫描\n" +
            f"  Scan ID: {scan_id}\n" +
            f"  Target: {target_name}\n" +
            f"  Workspace: {scan_workspace_dir}\n" +
            "="*60
        )
        
        # 参数验证
        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 不能为空")
        
        # Step 0: 创建工作目录
        from apps.scan.utils import setup_scan_directory
        directory_scan_dir = setup_scan_directory(scan_workspace_dir, 'directory_scan')
        
        # Step 1: 导出站点 URL（支持懒加载）
        sites_file, site_count = _export_site_urls(target_id, target_name, directory_scan_dir)
        
        if site_count == 0:
            logger.warning("目标下没有站点，跳过目录扫描")
            return {
                'success': True,
                'scan_id': scan_id,
                'target': target_name,
                'scan_workspace_dir': scan_workspace_dir,
                'sites_file': sites_file,
                'site_count': 0,
                'total_directories': 0,
                'processed_sites': 0,
                'failed_sites_count': 0,
                'executed_tasks': ['export_sites']
            }
        
        # Step 2: 工具配置信息
        logger.info("Step 2: 工具配置信息")
        tool_info = []
        for tool_name, tool_config in enabled_tools.items():
            mw = _get_max_workers(tool_config)
            tool_info.append(f"{tool_name}(max_workers={mw})")
        logger.info("✓ 启用工具: %s", ', '.join(tool_info))
        
        # Step 3: 并发执行扫描工具并实时保存结果
        logger.info("Step 3: 并发执行扫描工具并实时保存结果")
        total_directories, processed_sites, failed_sites = _run_scans_concurrently(
            enabled_tools=enabled_tools,
            sites_file=sites_file,
            directory_scan_dir=directory_scan_dir,
            scan_id=scan_id,
            target_id=target_id,
            site_count=site_count,
            target_name=target_name
        )
        
        # 检查是否所有站点都失败
        if processed_sites == 0 and site_count > 0:
            logger.warning("所有站点扫描均失败 - 总站点数: %d, 失败数: %d", site_count, len(failed_sites))
            # 不抛出异常，让扫描继续
        
        logger.info("="*60 + "\n✓ 目录扫描完成\n" + "="*60)
        
        return {
            'success': True,
            'scan_id': scan_id,
            'target': target_name,
            'scan_workspace_dir': scan_workspace_dir,
            'sites_file': sites_file,
            'site_count': site_count,
            'total_directories': total_directories,
            'processed_sites': processed_sites,
            'failed_sites_count': len(failed_sites),
            'executed_tasks': ['export_sites', 'run_and_stream_save_directories']
        }
        
    except Exception as e:
        logger.exception("目录扫描失败: %s", e)
        raise