"""
Copy Data Pipeline - Simple data copying without strategy execution
"""

from pathlib import Path
from typing import List, Any, Optional

from ..common.collection import DataCollection
from ..common.types import SyncAction, SyncStatus
from ..common.registry import DomainRegistry
from ..datasource import JsonlDataSource
from ..engine import StateMachineConfig, StateMachineRuntime, apply_state_to_node


class CopyDataPipeline:
    """
    通用数据复制 Pipeline，支持任意数据源之间的复制。
    
    流程：
    1. 从 source_datasource 加载所有数据到 source_collection
    2. 将每个节点复制到 target_collection，设置 action=CREATE
    3. 使用 target_datasource 保存所有节点
    
    用途：
    - 数据迁移 (JSONL → JSONL)
    - API 数据下载 (API → JSONL)
    - 测试数据准备
    - 初始化同步
    
    支持两种初始化方式：
    1. 传入数据源实例 (推荐): source_datasource + target_datasource
    2. 传入目录路径 (兼容): source_dir + target_dir (自动创建 JsonlDataSource)
    """
    
    def __init__(
        self,
        source_datasource: Any = None,
        target_datasource: Any = None,
        source_dir: Optional[Path | str] = None,
        target_dir: Optional[Path | str] = None,
        node_types: List[str] | None = None,
        load_order: List[str] | None = None,
        save_order: List[str] | None = None,
    ):
        """
        初始化 CopyDataPipeline
        
        Args:
            source_datasource: 源数据源实例 (ApiDataSource, JsonlDataSource 等)
            target_datasource: 目标数据源实例
            source_dir: 源数据目录 (兼容模式，会创建 JsonlDataSource)
            target_dir: 目标数据目录 (兼容模式，会创建 JsonlDataSource)
            node_types: 要复制的节点类型列表，None 表示复制所有注册的类型
            load_order: 数据加载顺序，None 表示使用 node_types 顺序
            save_order: 数据保存顺序，None 表示使用 node_types 顺序
        """
        self._supported_types: Optional[List[str]] = None

        # 数据源处理 - 优先使用实例，其次使用目录
        if source_datasource is not None:
            self.source_datasource = source_datasource
            self._owns_source = False
        elif source_dir is not None:
            from ..domain.registry import JSONL_HANDLERS
            self.source_datasource = JsonlDataSource(dir_path=Path(source_dir), read_only=True)
            # 自动注册所有 JSONL handlers，并收集支持的类型
            supported_types = []
            for handler_class in JSONL_HANDLERS:
                handler = handler_class(self.source_datasource)
                self.source_datasource.add_handler(handler)
                supported_types.append(handler.node_type)
            self._owns_source = True
            self._supported_types = supported_types
        else:
            raise ValueError("Must provide either source_datasource or source_dir")
        
        if target_datasource is not None:
            self.target_datasource = target_datasource
            self._owns_target = False
        elif target_dir is not None:
            from ..domain.registry import JSONL_HANDLERS
            self.target_datasource = JsonlDataSource(dir_path=Path(target_dir), read_only=False)
            # 自动注册所有 JSONL handlers
            for handler_class in JSONL_HANDLERS:
                self.target_datasource.add_handler(handler_class(self.target_datasource))
            self._owns_target = True
        else:
            raise ValueError("Must provide either target_datasource or target_dir")
        
        # 当使用 source_dir 创建时，如果 node_types 未指定，只使用支持的类型
        if node_types is None and self._supported_types is not None:
            self.node_types = self._supported_types
        else:
            self.node_types = node_types or DomainRegistry.get_registered_types()
        self.load_order = load_order or self.node_types
        self.save_order = save_order or self.node_types
        cfg_path = Path(__file__).resolve().parents[1] / "config" / "node_state_machine.yaml"
        self._sm_runtime = StateMachineRuntime(StateMachineConfig.load(cfg_path))
        
        # 统计信息
        self.stats = {
            "loaded": 0,
            "copied": 0,
            "saved": 0,
            "failed": 0,
        }
    
    async def run(self):
        """
        执行完整的复制流程
        
        Returns:
            dict: 统计信息 {"loaded": int, "copied": int, "saved": int, "failed": int}
        """
        print(f"[CopyDataPipeline] Starting...")
        print(f"  Source: {self.source_datasource.__class__.__name__}")
        print(f"  Target: {self.target_datasource.__class__.__name__}")
        print(f"  Node types: {self.node_types}")
        
        # Stage 1: 加载源数据
        print("\n[Stage 1] Loading source data...")
        source_collection = DataCollection(collection_id="source")
        await self._load_source_data(source_collection)
        
        # Stage 2: 复制节点
        print("\n[Stage 2] Copying nodes...")
        target_collection = DataCollection(collection_id="target")
        await self._copy_nodes(source_collection, target_collection)
        
        # Stage 3: 保存到目标
        print("\n[Stage 3] Saving to target...")
        await self._save_target_data(target_collection)
        
        # 打印统计信息
        self._print_summary()
        
        return self.stats
    
    async def _load_source_data(self, collection: DataCollection):
        """从源数据源加载所有节点"""
        # 如果是 JsonlDataSource 且没有注册 handlers，自动注册
        if isinstance(self.source_datasource, JsonlDataSource) and not self.source_datasource._handlers:
            self._auto_register_jsonl_handlers(self.source_datasource)
        
        self.source_datasource.set_collection(collection)
        await self.source_datasource.load_all(order=self.load_order)
        
        # 统计加载的节点数
        for node_type in self.node_types:
            count = len(collection.filter(node_type=node_type))
            if count > 0:
                print(f"  [LOAD] {node_type}: {count} nodes")
                self.stats["loaded"] += count
    
    async def _copy_nodes(
        self,
        source_collection: DataCollection,
        target_collection: DataCollection
    ):
        """
        复制所有节点到目标 collection，设置 action=CREATE, status=PENDING
        
        注意：
        - 保持相同的 node_id（用于追踪）
        - data_id 设为空字符串（待创建）
        - 复制 data 和 depend_ids
        """
        for node in source_collection.filter():  # filter() 不带参数返回所有节点
            # 获取节点类
            node_class = DomainRegistry.get_node_class(node.node_type)
            
            # 创建新节点
            new_node = node_class(
                node_id=node.node_id,      # 保持相同的 node_id
                data_id="",                # CREATE 节点的 data_id 初始为空
                data=node.data,            # 复制数据
                depend_ids=node.depend_ids.copy() if node.depend_ids else [],
            )

            apply_state_to_node(
                self._sm_runtime,
                new_node,
                state_id="S06",
                reason="copy pipeline init",
                fields={"binding_status", "action", "status"},
            )
            apply_state_to_node(
                self._sm_runtime,
                new_node,
                state_id="S06",
                reason="copy pipeline init",
                fields={"data_id"},
                data_id="",
            )
            
            await target_collection.add(new_node)
            self.stats["copied"] += 1
        
        print(f"  Copied {self.stats['copied']} nodes")
    
    async def _save_target_data(self, collection: DataCollection):
        """保存所有节点到目标数据源"""
        # 如果是 JsonlDataSource 且没有注册 handlers，自动注册
        if isinstance(self.target_datasource, JsonlDataSource) and not self.target_datasource._handlers:
            self._auto_register_jsonl_handlers(self.target_datasource)
        
        self.target_datasource.set_collection(collection)
        await self.target_datasource.sync_all(order=self.save_order)
        
        # 统计保存结果
        for node_type in self.node_types:
            nodes = collection.filter(node_type=node_type)
            success_count = sum(1 for n in nodes if n.status == SyncStatus.SUCCESS)
            failed_count = sum(1 for n in nodes if n.status == SyncStatus.FAILED)
            
            if success_count > 0 or failed_count > 0:
                print(f"  [SAVE] {node_type}: {success_count} success, {failed_count} failed")
            
            self.stats["saved"] += success_count
            self.stats["failed"] += failed_count
    
    def _auto_register_jsonl_handlers(self, datasource: JsonlDataSource):
        """自动为 JsonlDataSource 注册需要的 handlers（兼容模式）"""
        for node_type in self.node_types:
            if not DomainRegistry.is_registered(node_type):
                continue
            
            handler_class = DomainRegistry.get_jsonl_handler(node_type)
            if handler_class:
                handler = handler_class(datasource)
                datasource.add_handler(handler)
    
    def _print_summary(self):
        """打印执行摘要"""
        print("\n" + "=" * 60)
        print("[CopyDataPipeline] Execution Summary")
        print("=" * 60)
        print(f"  Loaded:  {self.stats['loaded']} nodes")
        print(f"  Copied:  {self.stats['copied']} nodes")
        print(f"  Saved:   {self.stats['saved']} nodes")
        print(f"  Failed:  {self.stats['failed']} nodes")
        print("=" * 60)
        
        if self.stats["failed"] > 0:
            print(f"\n⚠️  {self.stats['failed']} nodes failed to save")
        else:
            print("\n✅ All nodes saved successfully")
