tools/parameter_tools.py

from __future__ import annotations

from pathlib import Path

import pandas as pd

from langchain_core.tools import tool

from tools.config import PARAMETER_FILE_MAP, PARAMETERS_DIR


_FILENAME_TO_KEY: dict[str, str] = {
    "cluster_mapping.csv": "cluster_mapping",
    "market_summary_mapping.csv": "market_summary_mapping",
    "threshold_matrix.csv": "threshold_matrix",
    "country_region_mapping.csv": "country_region_mapping",
    "company_products.csv": "company_products",
}


def _resolve_parameter_keys(user_query: str) -> list[str]:
    query = (user_query or "").lower()
    selected: set[str] = set()
    for keyword, filename in PARAMETER_FILE_MAP.items():
        if keyword in query:
            key = _FILENAME_TO_KEY.get(filename)
            if key:
                selected.add(key)
    if not selected:
        selected.update(["cluster_mapping", "country_region_mapping"])
    return sorted(selected)


def _load_param_df(key: str) -> pd.DataFrame:
    """Load a parameter table from the SQLite DB; fall back to CSV if needed."""
    try:
        from db import load_param_from_db
        return load_param_from_db(key)
    except Exception:
        pass
    csv_name = {
        "cluster_mapping": "cluster_mapping.csv",
        "market_summary_mapping": "market_summary_mapping.csv",
        "threshold_matrix": "threshold_matrix.csv",
        "country_region_mapping": "country_region_mapping.csv",
        "company_products": "company_products.csv",
    }.get(key, "")
    if csv_name:
        path = Path(PARAMETERS_DIR) / csv_name
        if path.exists():
            return pd.read_csv(path, keep_default_na=False)
    return pd.DataFrame()


def _row_score(row_values: list[str], query_tokens: set[str]) -> int:
    joined = " ".join(row_values).lower()
    return sum(1 for token in query_tokens if token and token in joined)


@tool
def parameter_reader(user_query: str, max_rows_per_file: int = 25) -> list[dict]:
    """Read relevant parameter rows from the parameter tables (DB or CSV fallback)."""
    rows: list[dict] = []
    keys = _resolve_parameter_keys(user_query)
    query_tokens = {token for token in (user_query or "").lower().split() if len(token) > 2}

    for key in keys:
        try:
            df = _load_param_df(key)
        except Exception:
            continue
        if df.empty:
            continue

        records = df.to_dict(orient="records")
        if query_tokens:
            scored = []
            for record in records:
                score = _row_score([str(v) for v in record.values()], query_tokens)
                if score > 0:
                    scored.append((score, record))
            scored.sort(key=lambda item: item[0], reverse=True)
            selected_records = [record for _, record in scored[:max_rows_per_file]]
        else:
            selected_records = records[:max_rows_per_file]

        source_label = f"{key} (db)"
        for record in selected_records:
            rows.append({"_source": source_label, **record})

    return rows