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


class TransposeTool(BaseTool):
    def execute(self, inputs: Dict[str, pl.DataFrame]) -> Dict[str, pl.DataFrame]:
        df = inputs.get("Input", pl.DataFrame())
        if self.config is None or df.is_empty():
            return {"Output": df}

        keys = [
            f.attrib["name"]
            for f in self.config.findall("KeyFields/Field")
            if f.attrib["name"] in df.columns
        ]
        data_cols = [
            f.attrib["name"]
            for f in self.config.findall("DataFields/Field")
            if f.attrib["name"] in df.columns
        ]

        if not data_cols:
            # Default: all non-key columns become data columns
            data_cols = [c for c in df.columns if c not in keys]

        result = df.unpivot(
            on=data_cols,
            index=keys,
            variable_name="Name",
            value_name="Value",
        )
        return {"Output": result}