from __future__ import annotations

import logging
import os
from .util import markdown_color, markdown_terminal
from .tool_types import ToolOutput, ToolFile, extract_file_type

class Tool:
    logger = logging.getLogger("ADAM-TOOL")

    def __init_subclass__(cls, **kwargs):
        super().__init_subclass__()

        for key, value in kwargs.items():
            setattr(cls, key, value)
        if not hasattr(cls, "GPU"):
            cls.GPU = 0
        if not hasattr(cls, "CPU"):
            cls.CPU = 1
        if not hasattr(cls, "MEM_PER_CPU"):
            cls.MEM_PER_CPU = 4000
        if not hasattr(cls, "PARTITION"):
            cls.PARTITION = "gpu_3090"
        if not hasattr(cls, "CONDA_ENV"):
            cls.CONDA_ENV = "base"
        if not hasattr(cls, "calltype"):
            cls.calltype = "bash"
        if not hasattr(cls, "DISPLAY_NAME"):
            cls.DISPLAY_NAME = cls.__name__


    def inputShow(self, **kwargs):
        """
            AI决定使用该工具时调用，用于打印信息告知用户目前调用的工具
            此处的kwargs更改对后续过程无效

            :param dict args: 字典形式的AI产生的输入
            :param list[str] files: 此次用户上传的文件
            :param str message: 此次用户询问的消息
            :param str user: 用户昵称
            :param str task_id: 任务编号
        """
        if self.calltype == "bash":
            os.environ["ADAM_OUTPUT_RAW"] = "true"
            scripts = self.call(kwargs)
            os.environ["ADAM_OUTPUT_RAW"] = "false"

            scripts = '\n'.join(['bash -s << EOF', scripts, 'EOF'])
        else:
            scripts = f"""python {self.DISPLAY_NAME}.py"""
        return Tool.markdown_terminal(scripts, workdir=kwargs["task_id"],
                user=kwargs["user"], conda_env=self.CONDA_ENV)

    def resAlloc(self, kwargs):
        """
            提交工具任务时调用，用于获取分配的资源量
            此处的kwargs更改只有tool_data和tip对后续过程有效

            kwarg在inputShow的基础之上新增：
            :param any tool_data: 可序列化的任意数据
            :param str tip: 用来给用户显示的提示
            :returns: 一个字典，包括使用的GPU、CPU和MEM_PER_CPU（以MB单位的整数）
        """
        return {"CPU": self.CPU, "MEM_PER_CPU": self.MEM_PER_CPU,
                "GPU": self.GPU, "PARTITION": self.PARTITION}

    def call(self, kwargs):
        """
            实际执行脚本
            此处的kwargs更改只有tool_data，tip和token_usage对后续过程有效

            kwarg同resAlloc
            :returns: calltype=bash时，返回实际运行的bash脚本；
                calltype=python时，return调用值无效
        """
        raise NotImplementedError


    def monitor(self, kwargs):
        """
            任务运行时的监控命令，用于告知用户目前的进度
            只有在return为True时，此处的kwargs更改tool_data、tip和progress对后续过程有效

            kwargs在call的基础之上新增：
            :param int run_time: 以秒为单位的运行时间
            :param float progress: 任务进度，1为任务完成，0.5为50%，以此类推
        """
        pass

    def outputShow(self, kwargs):
        """
            任务结束时调用，用于打印信息告知用户工具结果

            kwargs在monitor的基础上新增：
            :param str stdout: stdout流的结果
            :param str stderr: stderr流的结果
            :param int exit_code: 退出码
            :returns: 一个给用户的文本, 一个包含显示文件的列表
        """
        stdout = kwargs["stdout"]
        if len(stdout) > 10000:
            stdout = stdout[:5000] + "\n...输出太长已省略...\n" + stdout[-5000:]
        
        stderr = kwargs["stderr"]
        if len(stderr) > 10000:
            stderr = stderr[:5000] + "\n...输出太长已省略...\n" + stderr[-5000:]
        return f'''- exit code: {kwargs['exit_code']}
- stdout:
```
{stdout}
```
- stderr:
```
{stderr}
```
''', []

    def summary(self, kwargs):
        """
            任务结束时调用，用于打印信息告知AI结果

            kwargs同outputShow
            :returns: 一个给AI的文本
        """
        stdout = kwargs["stdout"]
        if len(stdout) > 10000:
            stdout = stdout[:5000] + "\n...输出太长已省略...\n" + stdout[-5000:]
        stderr = kwargs["stderr"]
        if len(stderr) > 10000:
            stderr = stderr[:5000] + "\n...输出太长已省略...\n" + stderr[-5000:]
        return f'''- exit code: {kwargs['exit_code']}
- stdout:
```
{stdout}
```
- stderr:
```
{stderr}
```
'''

    def output(self, kwargs) -> ToolOutput:
        """Unified output method for DAG workflow data passing.

        Priority:
        1. Subclass overrides this method → use that directly.
        2. @tool_io outputs declared → auto-scan working directory by extension.
        3. Fallback → wrap legacy outputShow file list into ToolOutput.
        """
        text, legacy_files = self.outputShow(kwargs)
        smry = self.summary(kwargs)

        # Auto-match when @tool_io declares outputs
        io = getattr(self.__class__, '__tool_io__', None)
        if io and io.get('outputs'):
            matched = self._auto_match_outputs(kwargs, io['outputs'])
            if matched:
                return ToolOutput(files=matched, text=text, summary=smry)

        # Fallback: legacy outputShow bridge
        return ToolOutput(
            files={
                f"file_{i}": ToolFile(
                    type=os.path.splitext(f)[1].lstrip('.'),
                    description='',
                    path=f,
                )
                for i, f in enumerate(legacy_files or [])
            },
            text=text,
            summary=smry,
        )

    def _auto_match_outputs(self, kwargs, output_decl) -> dict[str, ToolFile]:
        """Scan working directory, match files to @tool_io output declarations by extension."""
        work_dir = kwargs.get('directory_path') or os.getcwd()
        if not os.path.isdir(work_dir):
            return {}

        # Collect all files recursively
        all_files = []
        for root, _, fnames in os.walk(work_dir):
            for fn in fnames:
                all_files.append(os.path.join(root, fn))

        # Group by extension
        by_ext: dict[str, list[str]] = {}
        for fpath in all_files:
            ext = os.path.splitext(fpath)[1].lstrip('.').lower()
            if ext:
                by_ext.setdefault(ext, []).append(fpath)

        # Match each declared output key to files
        result = {}
        claimed: set[str] = set()
        for key, annotated_type in output_decl.items():
            ft = extract_file_type(annotated_type)
            if not ft:
                continue
            ext = ft.ext.lstrip('.').lower()
            candidates = [f for f in by_ext.get(ext, []) if f not in claimed]
            if not candidates:
                continue
            candidates.sort()
            if len(candidates) == 1:
                result[key] = ToolFile(type=ext, description=ft.description, path=candidates[0])
                claimed.add(candidates[0])
            else:
                result[key] = ToolFile(type=ext, description=ft.description, paths=candidates)
                claimed.update(candidates)

        return result

    def fill_paths(self, paths: dict[str, str | list[str]]) -> dict[str, ToolFile]:
        """Merge runtime file paths with @tool_io output declarations.

        Values can be a single path ``str`` or a ``list[str]`` for
        multi-file outputs (e.g. N design variants).
        """
        io = getattr(self.__class__, '__tool_io__', {})
        output_decl = io.get('outputs', {})
        result = {}
        for key, val in paths.items():
            ft = extract_file_type(output_decl.get(key))
            if isinstance(val, list):
                first = val[0] if val else ''
                result[key] = ToolFile(
                    type=ft.ext.lstrip('.') if ft else os.path.splitext(first)[1].lstrip('.'),
                    description=ft.description if ft else '',
                    paths=val,
                )
            else:
                result[key] = ToolFile(
                    type=ft.ext.lstrip('.') if ft else os.path.splitext(val)[1].lstrip('.'),
                    description=ft.description if ft else '',
                    path=val,
                )
        return result

    def declare_outputs(self, files: dict[str, str | list[str]] | None = None, params: dict | None = None):
        """在 call() 末尾调用，声明本次执行的输出文件和参数值。

        写入 _outputs.json 到当前工作目录，Worker 执行完后读取此文件
        构建 structured_output，再由 dag_callback 按 io_contract 精确注册。

        Args:
            files: {port_name: file_path_or_list}，key 必须与 @tool_io(outputs=...) 一致。
                   路径为相对于工作目录的相对路径或绝对路径。
            params: {key: value}，非文件的标量输出（数值、字符串等）。

        Example::

            @tool_io(outputs={"structure": PDBFile, "scores": CSVFile})
            class MyTool(Tool):
                def call(self, kwargs):
                    # ... 计算 ...
                    self.declare_outputs(
                        files={"structure": "result.pdb", "scores": "scores.csv"},
                        params={"n_atoms": 1234, "rmsd": 0.85},
                    )
        """
        import json
        data: dict = {}
        if files:
            data["files"] = files
        if params:
            data["params"] = params
        with open("_outputs.json", "w", encoding="utf-8") as f:
            json.dump(data, f, ensure_ascii=False)

    @staticmethod
    def markdown_color(content, color):
        return markdown_color(content, color)

    @staticmethod
    def markdown_terminal(content, conda_env="base", user="Adam", workdir=""):
        return markdown_terminal(content, conda_env, user, workdir)
