# agent_app.py
from smolagents import CodeAgent, TransformersModel
from smolagents import tool
from sheet_tool import (
    fetch_sheet_as_df,
    create_pivot,
    summary_stats,
    plot_dataframe,
    df_to_csv_bytes,
)
import base64

# Initialize model
model = TransformersModel(model_id="HuggingFaceTB/SmolLM-135M-Instruct")

# -------------------------------
# ✅ TOOL DEFINITIONS
# -------------------------------

@tool
def load_sheet() -> dict:
    """Load Google Sheet into a dataframe and return a short summary (not the full sheet)."""
    df = fetch_sheet_as_df()
    if df.empty:
        return {"error": "Sheet is empty or not found."}
    return {
        "rows": len(df),
        "columns": list(df.columns),
        "head": df.head(5).to_dict(orient="records"),
    }


@tool
def pivot(index_cols: str, column_cols: str, value_cols: str, aggfunc: str = "sum") -> dict:
    """
    Create a pivot table from the Google Sheet.

    Args:
        index_cols (str): Comma-separated list of columns to use as the pivot table index.
        column_cols (str): Comma-separated list of columns to use as pivot table columns.
        value_cols (str): Comma-separated list of columns to aggregate.
        aggfunc (str, optional): Aggregation function to apply (e.g., 'sum', 'mean', 'count'). Defaults to 'sum'.
    """
    df = fetch_sheet_as_df()
    if df.empty:
        return {"error": "Sheet empty"}
    index = [c.strip() for c in index_cols.split(",")] if index_cols else []
    columns = [c.strip() for c in column_cols.split(",")] if column_cols else []
    values = [c.strip() for c in value_cols.split(",")] if value_cols else []
    pivot_df = create_pivot(df, index=index, columns=columns, values=values, aggfunc=aggfunc)
    csv_bytes = df_to_csv_bytes(pivot_df)
    return {
        "pivot_preview": pivot_df.head(10).to_dict(orient="records"),
        "csv_b64": base64.b64encode(csv_bytes).decode("utf-8")
    }



@tool
def stats() -> dict:
    """Generate summary statistics of the sheet."""
    df = fetch_sheet_as_df()
    if df.empty:
        return {"error": "Sheet empty"}
    s = summary_stats(df)
    return {"summary": s.to_dict()}


@tool
def plot(kind: str = "bar", x: str = None, y: str = None, title: str = None) -> dict:
    """
    Create a plot from the Google Sheet data.

    Args:
        kind (str): Type of chart to create. Example values: 'bar', 'line', 'pie', 'scatter'.
        x (str, optional): Column name to use for the X-axis. Example: 'Date'.
        y (str, optional): Comma-separated column names to use for Y-axis. Example: 'Sales,Profit'.
        title (str, optional): Chart title to display at the top.

    Returns:
        dict: A dictionary containing the base64-encoded plot image or an error message if the sheet is empty.
    """
    df = fetch_sheet_as_df()
    if df.empty:
        return {"error": "Sheet empty"}

    y_list = [c.strip() for c in y.split(",")] if y else None
    img_data_uri = plot_dataframe(df, kind=kind, x=x, y=y_list, title=title)
    return {"image": img_data_uri}


# -------------------------------
# ✅ AGENT CREATION
# -------------------------------
agent = CodeAgent(model=model, tools=[load_sheet, pivot, stats, plot], add_base_tools=True)


def ask_agent(nl_query: str) -> dict:
    """Send a natural-language query to the agent and return structured response."""
    try:
        resp = agent.run(nl_query)
        return {"text": str(resp)}
    except Exception as e:
        return {"error": str(e)}