import datetime
import os

from sodium import bench


def agent(query, schema, base_domain, primary_key, primary_key_vals, output_folder=None):
    """Run the SODIUM agent to fill a table via web exploration.

    Args:
        query:            Natural-language description of what to find.
        schema:           List of column names, or a pandas DataFrame whose
                          columns define the schema.
        base_domain:      Starting URL for web exploration (e.g. "https://example.gov/").
        primary_key:      Column name that identifies each row (e.g. "year").
        primary_key_vals: List of primary-key values to fill (e.g. [2020, 2021]).
        output_folder:    Directory to write output.csv, sources.csv, and cell logs.
                          Auto-generated under logs/ if omitted.

    Returns:
        pandas.DataFrame with the populated table.
    """
    if hasattr(schema, "columns"):
        schema = list(schema.columns)

    if output_folder is None:
        ts = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
        output_folder = os.path.join("logs", f"sodium_{ts}")
        os.makedirs(output_folder, exist_ok=True)

    from sodium._agent.agent import sodium_agent

    return sodium_agent(query, schema, primary_key, primary_key_vals, output_folder, base_domain)


__all__ = ["agent", "bench"]
