#!/usr/bin/env python3
"""
FileChannel 综合测试套件

测试 Sandai Operator SDK 的 FileChannel 核心功能
"""

import json
import time
import threading
import shutil
from pathlib import Path
import sys
import os

from sandai.operator import FileChannel

class TestResult:
    def __init__(self):
        self.passed = 0
        self.failed = 0
        self.tests = []
    
    def add_test(self, name, success, message=""):
        self.tests.append((name, success, message))
        if success:
            self.passed += 1
            print(f"✅ {name}")
        else:
            self.failed += 1
            print(f"❌ {name}: {message}")
    
    def summary(self):
        total = self.passed + self.failed
        print(f"\n{'='*50}")
        print(f"测试总结: {self.passed}/{total} 通过")
        if self.failed > 0:
            print(f"失败的测试:")
            for name, success, message in self.tests:
                if not success:
                    print(f"  - {name}: {message}")
        print('='*50)
        return self.failed == 0

def cleanup_test_files():
    """清理测试文件"""
    test_files = [
        "test_append_input.jsonl", "test_append_output.jsonl",
        "test_replace_input.jsonl", "test_replace_output.jsonl", 
        "test_polling_input.jsonl", "test_polling_output.jsonl",
        "test_batch_input.jsonl", "test_batch_output.jsonl",
        "test_vim_input.jsonl", "test_vim_output.jsonl"
    ]
    
    for file in test_files:
        if Path(file).exists():
            Path(file).unlink()

def test_file_append(result: TestResult):
    """测试文件追加场景"""
    test_input = "test_append_input.jsonl"
    test_output = "test_append_output.jsonl"
    
    try:
        # 清理并创建初始文件
        for file in [test_input, test_output]:
            if Path(file).exists():
                Path(file).unlink()
        
        with open(test_input, "w") as f:
            f.write('{"task_id": "append-001", "data": "initial"}\n')
        
        channel = FileChannel(test_input, test_output, debug=False)
        results = []
        
        def consumer():
            for i in range(3):
                try:
                    task = channel.pull(timeout=3.0)
                    results.append(task)
                    channel.push({"id": task["task_id"], "processed": True})
                except:
                    break
        
        consumer_thread = threading.Thread(target=consumer, daemon=True)
        consumer_thread.start()
        
        time.sleep(0.5)
        
        # 追加内容
        with open(test_input, "a") as f:
            f.write('{"task_id": "append-002", "data": "appended1"}\n')
            f.write('{"task_id": "append-003", "data": "appended2"}\n')
        
        consumer_thread.join(timeout=5)
        channel.close()
        
        success = len(results) == 3
        message = f"期望3个任务，实际{len(results)}个" if not success else ""
        result.add_test("文件追加测试", success, message)
        
    except Exception as e:
        result.add_test("文件追加测试", False, str(e))

def test_file_replacement(result: TestResult):
    """测试文件替换场景（模拟 vim 保存）"""
    test_input = "test_replace_input.jsonl"
    test_output = "test_replace_output.jsonl"
    
    try:
        # 清理并创建初始文件
        for file in [test_input, test_output]:
            if Path(file).exists():
                Path(file).unlink()
        
        with open(test_input, "w") as f:
            f.write('{"task_id": "replace-001", "data": "initial"}\n')
        
        channel = FileChannel(test_input, test_output, debug=False)
        results = []
        
        def consumer():
            for i in range(3):
                try:
                    task = channel.pull(timeout=4.0)
                    results.append(task)
                    channel.push({"id": task["task_id"], "processed": True})
                except:
                    break
        
        consumer_thread = threading.Thread(target=consumer, daemon=True)
        consumer_thread.start()
        
        time.sleep(0.5)
        
        # 模拟 vim 保存：创建临时文件然后替换
        temp_file = test_input + ".tmp"
        with open(temp_file, "w") as f:
            f.write('{"task_id": "replace-001", "data": "initial"}\n')
            f.write('{"task_id": "replace-002", "data": "replaced1"}\n')
            f.write('{"task_id": "replace-003", "data": "replaced2"}\n')
        
        shutil.move(temp_file, test_input)
        
        consumer_thread.join(timeout=6)
        channel.close()
        
        success = len(results) >= 2
        message = f"期望至少2个任务，实际{len(results)}个" if not success else ""
        result.add_test("文件替换测试", success, message)
        
    except Exception as e:
        result.add_test("文件替换测试", False, str(e))

def test_polling_fallback(result: TestResult):
    """测试轮询回退机制"""
    test_input = "test_polling_input.jsonl"
    test_output = "test_polling_output.jsonl"
    
    try:
        # 清理并创建初始文件
        for file in [test_input, test_output]:
            if Path(file).exists():
                Path(file).unlink()
        
        with open(test_input, "w") as f:
            f.write('{"task_id": "polling-001", "data": "initial"}\n')
        
        # 禁用文件监听，强制使用轮询
        channel = FileChannel(test_input, test_output, use_file_watcher=False, debug=False)
        results = []
        
        def consumer():
            for i in range(2):
                try:
                    task = channel.pull(timeout=3.0)
                    results.append(task)
                    channel.push({"id": task["task_id"], "processed": True})
                except:
                    break
        
        consumer_thread = threading.Thread(target=consumer, daemon=True)
        consumer_thread.start()
        
        time.sleep(1)
        
        # 追加内容
        with open(test_input, "a") as f:
            f.write('{"task_id": "polling-002", "data": "polled"}\n')
        
        consumer_thread.join(timeout=4)
        channel.close()
        
        success = len(results) == 2
        message = f"期望2个任务，实际{len(results)}个" if not success else ""
        result.add_test("轮询回退测试", success, message)
        
    except Exception as e:
        result.add_test("轮询回退测试", False, str(e))

def test_large_batch(result: TestResult):
    """测试大批量任务处理"""
    test_input = "test_batch_input.jsonl"
    test_output = "test_batch_output.jsonl"
    
    try:
        # 清理并创建初始文件
        for file in [test_input, test_output]:
            if Path(file).exists():
                Path(file).unlink()
        
        with open(test_input, "w") as f:
            f.write('{"task_id": "batch-001", "data": "initial"}\n')
        
        channel = FileChannel(test_input, test_output, debug=False)
        results = []
        
        def consumer():
            while len(results) < 10:
                try:
                    task = channel.pull(timeout=2.0)
                    results.append(task)
                    channel.push({"id": task["task_id"], "processed": True})
                except:
                    break
        
        consumer_thread = threading.Thread(target=consumer, daemon=True)
        consumer_thread.start()
        
        time.sleep(0.5)
        
        # 批量追加任务
        with open(test_input, "a") as f:
            for i in range(2, 11):
                f.write(f'{{"task_id": "batch-{i:03d}", "data": "batch_{i}"}}\n')
        
        consumer_thread.join(timeout=5)
        channel.close()
        
        success = len(results) == 10
        message = f"期望10个任务，实际{len(results)}个" if not success else ""
        result.add_test("大批量测试", success, message)
        
    except Exception as e:
        result.add_test("大批量测试", False, str(e))

def test_vim_simulation(result: TestResult):
    """测试 vim 编辑模拟"""
    test_input = "test_vim_input.jsonl"
    test_output = "test_vim_output.jsonl"
    
    try:
        # 清理并创建初始文件
        for file in [test_input, test_output]:
            if Path(file).exists():
                Path(file).unlink()
        
        initial_content = '{"task_id": "vim-001", "data": "initial"}\n'
        with open(test_input, "w") as f:
            f.write(initial_content)
        
        channel = FileChannel(test_input, test_output, debug=False)
        results = []
        
        def consumer():
            while len(results) < 3:
                try:
                    task = channel.pull(timeout=3.0)
                    results.append(task)
                    channel.push({"id": task["task_id"], "processed": True})
                except:
                    break
        
        consumer_thread = threading.Thread(target=consumer, daemon=True)
        consumer_thread.start()
        
        time.sleep(0.5)
        
        # 模拟 vim 编辑
        def simulate_vim_save(filepath, content):
            temp_file = filepath + ".tmp"
            with open(temp_file, 'w') as f:
                f.write(content)
            shutil.move(temp_file, filepath)
        
        # 第一次 vim 保存
        new_content = initial_content + '{"task_id": "vim-002", "data": "vim added"}\n'
        simulate_vim_save(test_input, new_content)
        
        time.sleep(1)
        
        # 第二次 vim 保存
        new_content += '{"task_id": "vim-003", "data": "vim added again"}\n'
        simulate_vim_save(test_input, new_content)
        
        consumer_thread.join(timeout=5)
        channel.close()
        
        success = len(results) >= 2
        message = f"期望至少2个任务，实际{len(results)}个" if not success else ""
        result.add_test("vim 模拟测试", success, message)
        
    except Exception as e:
        result.add_test("vim 模拟测试", False, str(e))

def main():
    """运行所有测试"""
    print("FileChannel 综合测试套件")
    print("=" * 50)
    
    result = TestResult()
    
    # 清理之前的测试文件
    cleanup_test_files()
    
    # 运行所有测试
    test_file_append(result)
    test_file_replacement(result)
    test_polling_fallback(result)
    test_large_batch(result)
    test_vim_simulation(result)
    
    # 清理测试文件
    cleanup_test_files()
    
    # 显示总结
    success = result.summary()
    
    if success:
        print("\n🎉 所有测试通过！")
        return 0
    else:
        print("\n❌ 部分测试失败！")
        return 1

if __name__ == "__main__":
    exit(main())