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


class FormulaTool(BaseTool):
    def execute(self, inputs: Dict[str, pl.DataFrame]) -> Dict[str, pl.DataFrame]:
        df = inputs.get("Input", pl.DataFrame())

        if self.config is None:
            return {"Output": df}

        for ff in self.config.findall("FormulaFields/FormulaField"):
            expr_text = ff.attrib.get("expression", "")
            field = ff.attrib.get("field", "")
            alteryx_type = ff.attrib.get("type", "V_WString")
            size = ff.attrib.get("size")
            dtype = self.ctx.type_mapper.map(alteryx_type, size)

            if not field:
                continue

            try:
                series = self.ctx.transpiler.eval_series(df, expr_text, field, dtype)
            except Exception as e:
                raise RuntimeError(
                    f"Formula field {field!r} expression {expr_text!r} failed: {e}"
                ) from e

            df = df.with_columns(series.alias(field))

        return {"Output": df}