"""Units Domain - JSONL Handler"""
import json
import re
from typing import List

from ...datasource import BaseJsonlHandler, JsonlDataSource
from .sync_node import UnitsSyncNode, SyncNode
from schemas.project_extensions.generator_unit import GeneratorUnitProject


class UnitsJsonlHandler(BaseJsonlHandler):
    """机组 JSONL Handler"""
    def __init__(self, datasource: JsonlDataSource):
        super().__init__(datasource, "units", UnitsSyncNode, GeneratorUnitProject)

    async def load(self) -> List['SyncNode']:
        """加载机组数据（支持 generator_unit_N.jsonl 与 units_N.jsonl）"""
        nodes: List[SyncNode] = []
        seen_ids: set[str] = set()

        if not self.datasource.dir_path.exists():
            return nodes

        # 查找所有匹配的文件：generator_unit_N.jsonl 或 units_N.jsonl
        patterns = (
            re.compile(r"^generator_unit_\d+\.jsonl$"),
            re.compile(r"^units_\d+\.jsonl$"),
        )

        for file_path in self.datasource.dir_path.iterdir():
            if not any(pattern.match(file_path.name) for pattern in patterns):
                continue

            # 读取文件
            with open(file_path, "r", encoding="utf-8") as f:
                for line_num, line in enumerate(f, 1):
                    line = line.strip()
                    if not line:
                        continue

                    try:
                        item = json.loads(line)
                        item_id = self.extract_id(item)
                        if item_id and item_id in seen_ids:
                            continue

                        # 创建节点
                        node = self.create_node(item)
                        nodes.append(node)

                        # 同时加载到缓存
                        if item_id:
                            seen_ids.add(item_id)
                            bucket = self.datasource._data.setdefault(self.node_type, {})
                            bucket[item_id] = item

                    except json.JSONDecodeError as e:
                        print(f"Warning: Invalid JSON in {file_path.name}:{line_num}: {e}")

        return nodes