import sys
import json
import importlib.util
import re
from pathlib import Path
from typing import Any, List, Callable
from functools import cache

from pydantic import BaseModel

from alayaflow.workflow.runnable import StateGraphRunnableWorkflow
from alayaflow.workflow.workflow_info import WorkflowInfo
from alayaflow.workflow.runnable import BaseRunnableWorkflow


class WorkflowLoader:
    def __init__(self, workflow_dir: Path | str, workflow_id: str, version: str) -> None:
        self.workflow_dir = Path(workflow_dir) if isinstance(workflow_dir, str) else workflow_dir
    
    @cache
    def load_info(self) -> WorkflowInfo:
        if not self.workflow_dir.exists():
            raise FileNotFoundError(f"工作流目录 {self.workflow_dir} 不存在")
            
        metadata_file = self.workflow_dir / "metadata.json"
        if not metadata_file.exists():
            raise FileNotFoundError(f"工作流目录 {self.workflow_dir} 中没有 metadata.json 文件")
        
        try:
            with open(metadata_file, 'r', encoding='utf-8') as f:
                metadata = json.load(f)
        except json.JSONDecodeError as e:
            raise ValueError(f"解析 metadata.json 失败: {e}")
        except Exception as e:
            raise IOError(f"读取 metadata.json 失败: {e}")
        
        required_fields = ["id", "version", "entry_file", "entry_point"]
        missing_fields = [field for field in required_fields if field not in metadata]
        
        if missing_fields:
            raise ValueError(
                f"工作流目录 {self.workflow_dir} 的 metadata.json 中缺少必需字段: {', '.join(missing_fields)}"
            )
        
        return WorkflowInfo(
            id=metadata["id"],
            name=metadata.get("name", metadata["id"]),
            description=metadata.get("description", ""),
            version=metadata.get("version"),
            tags=metadata.get("tags", []),
            entry_file=metadata["entry_file"],
            entry_point=metadata["entry_point"],
            wf_dir=self.workflow_dir,
        )
    
    @cache
    def load_requirements(self) -> List[str]:
        requirements_file = self.workflow_dir / "requirements.txt"
        if not requirements_file.exists():
            return []
        
        requirements: List[str] = []
        try:
            with open(requirements_file, 'r', encoding='utf-8') as f:
                for line in f:
                    line = line.strip()
                    if line and not line.startswith("#"):
                        requirements.append(line)
        except Exception as e:
            raise IOError(f"读取 requirements.txt 失败: {e}")
        
        return requirements

    def _clear_module_cache(self, module_name: str) -> None:
        """清理模块及其子模块缓存（用于 force_reload）"""
        prefix = module_name + "."
        for key in list(sys.modules.keys()):
            if key == module_name or key.startswith(prefix):
                del sys.modules[key]

    def _load_runnable_data(self, force_reload: bool = False) -> tuple[Callable[..., BaseModel], Callable[..., Any]]:
        info = self.load_info()

        entry_filepath = info.wf_dir / info.entry_file
        if not entry_filepath.exists():
            raise FileNotFoundError(
                f"工作流 {info.id} 的入口文件不存在: {entry_filepath}"
            )

        module_name = self._sanitize_module_name(info.id, info.version)

        if force_reload:
            self._clear_module_cache(module_name)
        module = sys.modules.get(module_name)
        if module is None:
            module = self._load_module(module_name, entry_filepath)

        create_graph_func = self._resolve_entry_point(module, info.entry_point, info.id)
        if not callable(create_graph_func):
            raise TypeError(
                f"工作流 {info.id} 的入口点 '{info.entry_point}' 不是可调用对象，"
                f"而是 {type(create_graph_func)}"
            )

        # TODO: Support loading input schema
        get_input_schema_func = None

        return get_input_schema_func, create_graph_func

    def load_workflow(self, init_args: dict = None, force_reload: bool = False) -> BaseRunnableWorkflow:
        """加载工作流，返回包含 info 和 creator 的加载结果"""
        if init_args is None:
            init_args = {}

        info = self.load_info()

        # now we only support StateGraph
        get_input_schema_func, create_graph_func = self._load_runnable_data(force_reload)
        graph = create_graph_func(init_args)
        runnable = StateGraphRunnableWorkflow(info=info, graph=graph)

        return runnable
    
    @classmethod
    def _sanitize_module_name(cls, wf_id: str, version: str) -> str:
        safe_version = version.replace(".", "_")
        safe_wf_id = re.sub(r'[^0-9a-zA-Z_]', '_', wf_id)
        return f"workflow.{safe_wf_id}.{safe_version}"
        
    @classmethod
    def _load_module(cls, module_name: str, filepath: Path) -> Any:
        """加载 Python 模块，支持相对导入"""
        try:
            # 工作流目录（包含 workflow.py 的目录）
            workflow_dir = filepath.parent
            
            # 确保所有父模块都存在（用于支持相对导入）
            # 例如：workflow.autotable.1_1_0 需要 workflow 和 workflow.autotable 存在
            parts = module_name.split('.')
            for i in range(1, len(parts)):
                parent_name = '.'.join(parts[:i])
                if parent_name not in sys.modules:
                    # 创建父模块占位符
                    parent_module = type(sys)('module')
                    parent_module.__name__ = parent_name
                    parent_module.__package__ = parent_name
                    # 对于包目录，设置 __path__ 为工作流目录
                    # 这样相对导入可以找到同目录下的模块
                    parent_module.__path__ = [str(workflow_dir)]
                    sys.modules[parent_name] = parent_module
            
            # 使用 spec_from_file_location 加载模块
            # submodule_search_locations 参数使相对导入能够工作
            spec = importlib.util.spec_from_file_location(
                module_name,
                filepath,
                submodule_search_locations=[str(workflow_dir)]
            )
            if spec is None:
                raise ImportError(f"无法为文件 {filepath} 创建模块规范")
            
            if spec.loader is None:
                raise ImportError(f"文件 {filepath} 没有加载器")
            
            module = importlib.util.module_from_spec(spec)
            # 设置 __package__ 属性以支持相对导入
            # 对于包内的模块，__package__ 应该是包含它的包名
            if '.' in module_name:
                module.__package__ = module_name.rsplit('.', 1)[0]
            else:
                module.__package__ = module_name
            
            # 注册模块到 sys.modules（必须在 exec_module 之前）
            sys.modules[module_name] = module
            spec.loader.exec_module(module)
            
            return module
        except Exception as e:
            # 如果加载失败，清理 sys.modules
            if module_name in sys.modules:
                del sys.modules[module_name]
            raise ImportError(f"加载工作流模块失败: {e}") from e
    
    @classmethod
    def _resolve_entry_point(cls, module: Any, entry_point: str, workflow_id: str) -> Any:
        parts = entry_point.split('.')
        obj = module
        
        for i, part in enumerate(parts):
            if not hasattr(obj, part):
                path_so_far = '.'.join(parts[:i+1])
                raise AttributeError(
                    f"工作流 {workflow_id} 的模块中没有找到入口点 '{path_so_far}'"
                )
            obj = getattr(obj, part)
        
        return obj
