import os
import json
import time

from .config import ADAM_TASK_DIR


class _StatesManager:
    """内部状态管理类"""
    _tool_name = "tool"

    def _get_states_file(self) -> str:
        """获取 states.json 文件路径"""
        # 1. 优先使用环境变量
        task_dir = ADAM_TASK_DIR
        if task_dir:
            return os.path.join(task_dir, ".slurm", "states.json")

        # 2. 尝试当前目录的 .slurm
        current_dir = os.getcwd()
        local_states = os.path.join(current_dir, ".slurm", "states.json")
        if os.path.exists(local_states):
            return local_states

        # 3. 在当前目录创建
        return local_states

    def _read(self) -> dict:
        """读取 states.json，保留所有来源的数据"""
        states_file = self._get_states_file()
        if not os.path.exists(states_file):
            return {
                "files": {"updated_at": None, "sources": {}},
                "states": {"source": None, "updated_at": None, "data": {}}
            }
        try:
            with open(states_file, 'r', encoding='utf-8') as f:
                return json.load(f)
        except (json.JSONDecodeError, IOError):
            return {
                "files": {"updated_at": None, "sources": {}},
                "states": {"source": None, "updated_at": None, "data": {}}
            }

    def _write(self, data: dict):
        """写入 states.json"""
        states_file = self._get_states_file()
        os.makedirs(os.path.dirname(states_file), exist_ok=True)
        with open(states_file, 'w', encoding='utf-8') as f:
            json.dump(data, f, ensure_ascii=False, indent=2)

    def _set_nested(self, data: dict, key: str, value):
        """设置嵌套值，支持 dot notation"""
        keys = key.split('.')
        current = data
        for k in keys[:-1]:
            current = current.setdefault(k, {})
        current[keys[-1]] = value

    def _get_nested(self, data: dict, key: str, default=None):
        """获取嵌套值，支持 dot notation"""
        keys = key.split('.')
        current = data
        for k in keys:
            if isinstance(current, dict):
                current = current.get(k)
            else:
                return default
            if current is None:
                return default
        return current

    def set(self, key: str, value):
        """
        设置状态

        特殊 key:
            "files" - 记录文件列表，会自动补充到 files.sources[tool_name]
            其他 - 记录到 states.data
        """
        data = self._read()

        if key == "files":
            data["files"]["sources"][self._tool_name] = {
                "updated_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
                "items": value
            }
        else:
            self._set_nested(data["states"]["data"], key, value)

        self._write(data)

    def get(self, key: str):
        """
        获取状态

        Returns:
            状态值，不存在返回 None
        """
        data = self._read()

        if key == "files":
            # 合并所有来源的文件列表
            merged = {}
            for source, content in data["files"]["sources"].items():
                for item in content.get("items", []):
                    path = item.get("path")
                    if not path:
                        continue
                    mtime = item.get("mtime", 0)
                    if path not in merged or mtime > merged[path]["mtime"]:
                        merged[path] = item
            result = list(merged.values())
            # 按 mtime 降序排序
            result.sort(key=lambda x: x.get("mtime", 0), reverse=True)
            return result

        return self._get_nested(data["states"]["data"], key)

    def cleanup(self):
        """清理当前工具的所有记录"""
        data = self._read()
        data["files"]["sources"].pop(self._tool_name, None)
        self._write(data)

    def trackPath(self, path: str, max_items: int = 30):
        """
        追踪文件/目录，自动检测文件信息并记录

        Args:
            path: 文件或目录路径
            max_items: 最大记录数量，先进先出
        """
        if not os.path.exists(path):
            return

        # 检测文件信息
        try:
            is_dir = os.path.isdir(path)
            mtime = int(os.path.getmtime(path))
            # 从环境变量获取 task_dir
            task_dir = os.getenv('ADAM_TASK_DIR')
            if task_dir:
                rel_path = os.path.relpath(path, task_dir)
            else:
                rel_path = path
        except (OSError, ValueError):
            return

        # 获取当前文件列表
        data = self._read()
        source = self._tool_name
        items = data["files"]["sources"].get(source, {}).get("items", [])

        # 检查是否已存在相同路径，存在则移除（后面会加到末尾，保持最新追踪的在最后）
        items = [item for item in items if item.get("path") != rel_path]

        # 添加到末尾（最新的位置）
        items.append({"path": rel_path, "is_dir": is_dir, "mtime": mtime})

        # 先进先出，保留最多 max_items 条
        if len(items) > max_items:
            items = items[-max_items:]

        # 写回
        data["files"]["sources"][source] = {
            "updated_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
            "items": items
        }
        self._write(data)


# 全局实例
_states = _StatesManager()


def setState(key: str, value):
    """设置状态"""
    return _states.set(key, value)


def getState(key: str):
    """获取状态"""
    return _states.get(key)


def trackPath(path: str, max_items: int = 30):
    """
    追踪文件/目录，自动检测文件信息并记录

    Args:
        path: 文件或目录路径
        max_items: 最大记录数量，先进先出，默认 30
    """
    return _states.trackPath(path, max_items)
