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


class RecordIDTool(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}

        field = self._cfg("Field", "RecordID") or "RecordID"
        start_str = self._cfg("StartValue", "1") or "1"
        start = int(start_str)
        alteryx_type = self._cfg("FieldType", "Int32") or "Int32"
        dtype = self.ctx.type_mapper.map(alteryx_type)

        df = df.with_row_index(name=field, offset=start)
        df = df.with_columns(pl.col(field).cast(dtype))

        # Move the ID column to first position
        cols = [field] + [c for c in df.columns if c != field]
        return {"Output": df.select(cols)}