from __future__ import annotations
from typing import Dict
import polars as pl
from tools.base import BaseTool


class UnionTool(BaseTool):
    def execute(self, inputs: Dict[str, pl.DataFrame]) -> Dict[str, pl.DataFrame]:
        if not inputs:
            return {"Output": pl.DataFrame()}

        mode = self._cfg("Mode", "Auto") or "Auto"

        # Sort inputs: Input, Input1, Input2... (Alteryx uses 1-based suffixes)
        def _sort_key(k: str) -> int:
            suffix = k.replace("Input", "").lstrip("#")
            return int(suffix) if suffix.isdigit() else 0

        sorted_keys = sorted(inputs.keys(), key=_sort_key)
        dfs = [inputs[k] for k in sorted_keys if inputs[k] is not None]

        if not dfs:
            return {"Output": pl.DataFrame()}

        if mode == "ByPosition":
            names = dfs[0].columns
            renamed = []
            for df in dfs:
                if len(df.columns) == len(names):
                    renamed.append(df.rename(dict(zip(df.columns, names))))
                else:
                    renamed.append(df)
            result = pl.concat(renamed, how="diagonal_relaxed")
        else:
            result = pl.concat(dfs, how="diagonal_relaxed")

        return {"Output": result}