#!/usr/bin/env python3
"""
ProcessingContext 工作目录功能测试

测试 Sandai Operator SDK 的 ProcessingContext 核心功能，包括：
- 工作目录创建和管理
- 任务隔离
- 目录清理
- 错误处理
"""

import os
import sys
import tempfile
import json
from pathlib import Path
from typing import List, Generator

from sandai.operator import ProcessingContext, TaskInput, TaskOutput, InputArtifact
from pydantic import BaseModel


class TestOptions(BaseModel):
    """测试用的选项类"""
    test_param: str = "default"


class TestResults(BaseModel):
    """测试用的结果类"""
    files_created: int = 0
    workdir_path: str = ""


def test_processing_context_basic():
    """测试 ProcessingContext 基本功能"""
    print("🧪 测试 ProcessingContext 基本功能")
    
    with tempfile.TemporaryDirectory() as temp_dir:
        context = ProcessingContext(temp_dir)
        
        # 测试获取工作目录
        task_id = "test_task_001"
        workdir = context.get_task_workdir(task_id)
        
        # 验证路径格式
        expected_path = str(Path(temp_dir) / task_id / "workdir")
        assert workdir == expected_path, f"工作目录路径不正确: {workdir} != {expected_path}"
        
        # 验证目录存在
        assert Path(workdir).exists(), f"工作目录不存在: {workdir}"
        assert Path(workdir).is_dir(), f"工作目录不是目录: {workdir}"
        
        print(f"✅ 基本功能测试通过")


def test_processing_context_isolation():
    """测试任务隔离功能"""
    print("🧪 测试任务隔离功能")
    
    with tempfile.TemporaryDirectory() as temp_dir:
        context = ProcessingContext(temp_dir)
        
        # 创建多个任务的工作目录
        task_ids = ["task_001", "task_002", "task_003"]
        workdirs = {}
        
        for task_id in task_ids:
            workdir = context.get_task_workdir(task_id)
            workdirs[task_id] = workdir
            
            # 在每个工作目录中创建测试文件
            test_file = Path(workdir) / f"{task_id}_test.txt"
            test_file.write_text(f"测试文件 for {task_id}")
        
        # 验证所有工作目录都不同
        unique_workdirs = set(workdirs.values())
        assert len(unique_workdirs) == len(task_ids), "工作目录应该都不相同"
        
        # 验证每个工作目录都存在且包含正确的文件
        for task_id, workdir in workdirs.items():
            test_file = Path(workdir) / f"{task_id}_test.txt"
            assert test_file.exists(), f"测试文件不存在: {test_file}"
            content = test_file.read_text()
            assert content == f"测试文件 for {task_id}", f"文件内容不正确: {content}"
        
        print(f"✅ 任务隔离测试通过")


def test_processing_context_file_creation():
    """测试工作目录中的文件创建功能"""
    print("🧪 测试工作目录文件创建功能")
    
    with tempfile.TemporaryDirectory() as temp_dir:
        context = ProcessingContext(temp_dir)
        
        # 创建工作目录并添加文件
        task_id = "file_test_001"
        workdir = context.get_task_workdir(task_id)
        
        # 创建多层目录结构和文件
        sub_dir = Path(workdir) / "subdir"
        sub_dir.mkdir()
        
        files_to_create = [
            Path(workdir) / "test1.txt",
            Path(workdir) / "test2.json",
            sub_dir / "nested_file.txt"
        ]
        
        for file_path in files_to_create:
            file_path.write_text(f"测试内容: {file_path.name}")
        
        # 验证文件都存在
        task_dir = Path(temp_dir) / task_id
        assert task_dir.exists(), "任务目录应该存在"
        for file_path in files_to_create:
            assert file_path.exists(), f"文件应该存在: {file_path}"
        
        # 验证目录结构正确
        assert (Path(workdir) / "test1.txt").read_text() == "测试内容: test1.txt"
        assert (sub_dir / "nested_file.txt").read_text() == "测试内容: nested_file.txt"
        
        print(f"✅ 文件创建功能测试通过")


def test_processing_context_multiple_calls():
    """测试多次调用 get_task_workdir 的行为"""
    print("🧪 测试多次调用 get_task_workdir")
    
    with tempfile.TemporaryDirectory() as temp_dir:
        context = ProcessingContext(temp_dir)
        
        task_id = "multi_call_test"
        
        # 多次调用应该返回相同的路径
        workdir1 = context.get_task_workdir(task_id)
        workdir2 = context.get_task_workdir(task_id)
        workdir3 = context.get_task_workdir(task_id)
        
        assert workdir1 == workdir2 == workdir3, "多次调用应该返回相同路径"
        
        # 目录应该只创建一次
        assert Path(workdir1).exists(), "工作目录应该存在"
        
        print(f"✅ 多次调用测试通过")


def test_processing_context_in_batch_processing():
    """测试在批处理中使用 ProcessingContext"""
    print("🧪 测试批处理中的 ProcessingContext 使用")
    
    def mock_batch_processor(
        batch_inputs: List[TaskInput[TestOptions]], 
        operator_config: dict,
        context: ProcessingContext
    ) -> Generator[TaskOutput[TestResults], None, None]:
        """模拟批处理函数"""
        
        for task_input in batch_inputs:
            try:
                # 获取工作目录
                workdir = context.get_task_workdir(task_input.task_id)
                
                # 在工作目录中创建文件
                files_created = 0
                
                # 创建配置文件
                config_file = Path(workdir) / "config.json"
                config_data = {
                    "task_id": task_input.task_id,
                    "test_param": task_input.options.test_param
                }
                config_file.write_text(json.dumps(config_data, indent=2))
                files_created += 1
                
                # 创建处理日志
                log_file = Path(workdir) / "processing.log"
                log_file.write_text(f"处理任务: {task_input.task_id}\n")
                files_created += 1
                
                # 模拟处理输入文件
                for i, artifact in enumerate(task_input.artifacts or []):
                    temp_file = Path(workdir) / f"processed_{i}.tmp"
                    temp_file.write_text(f"处理结果: {artifact.uri}")
                    files_created += 1
                
                # 返回结果
                yield TaskOutput[TestResults](
                    task_id=task_input.task_id,
                    results=TestResults(
                        files_created=files_created,
                        workdir_path=workdir
                    ),
                    status="success"
                )
                
            except Exception as e:
                yield TaskOutput[TestResults](
                    task_id=task_input.task_id,
                    error={
                        "error_type": "ProcessingError",
                        "error_message": str(e)
                    },
                    status="error"
                )
    
    with tempfile.TemporaryDirectory() as temp_dir:
        context = ProcessingContext(temp_dir)
        
        # 创建测试输入
        batch_inputs = [
            TaskInput[TestOptions](
                task_id="batch_test_001",
                options=TestOptions(test_param="param1"),
                artifacts=[
                    InputArtifact(uri="file1.txt"),
                    InputArtifact(uri="file2.txt")
                ]
            ),
            TaskInput[TestOptions](
                task_id="batch_test_002",
                options=TestOptions(test_param="param2"),
                artifacts=[
                    InputArtifact(uri="file3.txt")
                ]
            )
        ]
        
        # 处理批次
        results = list(mock_batch_processor(batch_inputs, {}, context))
        
        # 验证结果
        assert len(results) == 2, f"应该有2个结果，实际有{len(results)}个"
        
        for result in results:
            assert result.status == "success", f"任务应该成功: {result.task_id}"
            assert result.results is not None, f"结果不应该为空: {result.task_id}"
            assert result.results.files_created > 0, f"应该创建了文件: {result.task_id}"
            
            # 验证工作目录存在且包含文件
            workdir = Path(result.results.workdir_path)
            task_dir = workdir.parent  # workdir 的父目录是任务目录
            assert task_dir.exists(), f"任务目录应该存在: {task_dir}"
            assert workdir.exists(), f"工作目录应该存在: {workdir}"
        
        print(f"✅ 批处理测试通过")


def run_all_processing_context_tests():
    """运行所有 ProcessingContext 测试"""
    print("🚀 开始 ProcessingContext 测试套件\n")
    
    tests = [
        test_processing_context_basic,
        test_processing_context_isolation,
        test_processing_context_file_creation,
        test_processing_context_multiple_calls,
        test_processing_context_in_batch_processing
    ]
    
    passed = 0
    failed = 0
    
    for test_func in tests:
        try:
            test_func()
            passed += 1
            print()
        except Exception as e:
            print(f"❌ 测试失败: {test_func.__name__}")
            print(f"   错误: {e}")
            failed += 1
            print()
    
    print(f"📊 测试结果: {passed} 通过, {failed} 失败")
    
    if failed == 0:
        print("🎉 所有 ProcessingContext 测试通过！")
        return True
    else:
        print("💥 部分测试失败")
        return False


if __name__ == "__main__":
    success = run_all_processing_context_tests()
    sys.exit(0 if success else 1)