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


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

        fmt = self._cfg("Format", "%Y-%m-%d") or "%Y-%m-%d"
        input_field = self._cfg("InputField", "") or ""
        output_field = self._cfg("OutputField", input_field) or input_field
        input_type = self._cfg("InputType", "String") or "String"
        output_type = self._cfg("OutputType", "Date") or "Date"

        if not input_field or input_field not in df.columns:
            return {"Output": df}

        # Convert Alteryx format to Python/Polars strptime format (they match mostly)
        col = pl.col(input_field)

        if input_type == "String":
            if output_type == "Date":
                result = col.str.to_date(fmt, strict=False)
            elif output_type == "DateTime":
                result = col.str.to_datetime(fmt, strict=False)
            elif output_type == "Time":
                result = col.str.to_time(fmt, strict=False)
            else:
                result = col.str.to_datetime(fmt, strict=False)
        else:
            # Date/DateTime → String
            if input_type == "Date":
                result = col.dt.strftime(fmt)
            elif input_type == "DateTime":
                result = col.dt.strftime(fmt)
            else:
                result = col.cast(pl.String)

        df = df.with_columns(result.alias(output_field))
        return {"Output": df}