"""Chat assistant: LLM-backed via HF Inference, with a rule-based fallback.

If the HF_TOKEN environment variable is set (a Space secret), replies come
from a hosted chat model (CHAT_MODEL env var, default Gemma 3 27B). On any
failure — missing token, exhausted quota, provider error — the assistant
falls back to the built-in keyword knowledge base, so the tab always works.
"""

from __future__ import annotations

import os
import re

CHAT_MODEL = os.environ.get("CHAT_MODEL", "google/gemma-3-27b-it")

WELCOME = (
    "Hi! I'm your ML analysis assistant. Ask me anything about machine learning "
    "models (e.g. *'What is a random forest?'*, *'When should I use logistic "
    "regression?'*), or upload a dataset in the **Data & Preprocessing** tab.\n\n"
    "**So — what do you want to analyse today?**"
)

SYSTEM_PROMPT = (
    "You are the built-in assistant of 'ML Data Analysis Studio', a Gradio app for "
    "no-code machine learning analysis. The app has four tabs:\n"
    "1. ML Assistant — this chat.\n"
    "2. Data & Preprocessing — upload CSV/TSV/Excel/JSON/Parquet (delimiter/encoding "
    "auto-detected), automatic profiling, then cleaning: mean/median imputation or "
    "row-dropping, duplicate removal, IQR outlier clipping, one-hot encoding, and "
    "standard or min-max scaling. Placeholders ('N/A', '?'), numeric-looking text "
    "('$1,234', '45%'), and empty/constant columns are handled automatically.\n"
    "3. Model Training — regression (Linear, Decision Tree, Random Forest, Gradient "
    "Boosting, SVR, KNN) or classification (Logistic, Decision Tree, Random Forest, "
    "Gradient Boosting, SVC, KNN), with metrics, plots, feature importance, and the "
    "equivalent scikit-learn code shown in a dropdown.\n"
    "4. Reports & Downloads — PDF/HTML report and cleaned CSV export.\n\n"
    "Answer questions about machine learning, statistics, and data analysis clearly "
    "and concisely (a few short paragraphs at most, markdown allowed). When relevant, "
    "point the user to the right tab and to the models this app actually offers. "
    "If asked about something unrelated to ML or data analysis, answer briefly and "
    "steer the conversation back to their data."
)


def _llm_reply(message: str, history: list | None, uploaded_file_name: str | None) -> str:
    from huggingface_hub import InferenceClient

    client = InferenceClient(model=CHAT_MODEL, token=os.environ.get("HF_TOKEN"))
    system = SYSTEM_PROMPT
    if uploaded_file_name:
        system += (
            f"\n\nThe user has already uploaded the dataset '{uploaded_file_name}' "
            "in the Data & Preprocessing tab."
        )
    messages = [{"role": "system", "content": system}]
    for m in (history or [])[-8:]:
        if (
            isinstance(m, dict)
            and m.get("role") in ("user", "assistant")
            and isinstance(m.get("content"), str)
        ):
            messages.append({"role": m["role"], "content": m["content"]})
    messages.append({"role": "user", "content": message})
    out = client.chat_completion(messages=messages, max_tokens=512, temperature=0.4)
    reply = out.choices[0].message.content
    if not reply or not reply.strip():
        raise ValueError("empty LLM reply")
    return reply.strip()


# ------------------------------------------------- rule-based fallback ----

KNOWLEDGE_BASE = {
    ("linear regression",): (
        "**Linear Regression** fits a straight line (or hyperplane) that minimizes the "
        "squared error between predictions and actual values. Use it when the target is "
        "continuous and roughly linearly related to the features. Pros: fast, interpretable "
        "coefficients. Cons: can't capture non-linear patterns, sensitive to outliers."
    ),
    ("logistic regression",): (
        "**Logistic Regression** is a classification model that predicts class "
        "probabilities via the sigmoid function. Great baseline for binary classification. "
        "Pros: fast, well-calibrated probabilities, interpretable. Cons: linear decision "
        "boundary only."
    ),
    ("random forest",): (
        "**Random Forest** builds many decision trees on bootstrapped samples with random "
        "feature subsets, then averages (regression) or votes (classification). Pros: "
        "handles non-linearity, robust to outliers/overfitting, gives feature importance. "
        "Cons: slower and less interpretable than a single tree."
    ),
    ("decision tree",): (
        "**Decision Trees** split data by feature thresholds to form a flowchart of rules. "
        "Pros: highly interpretable, no scaling needed, handles mixed data. Cons: prone to "
        "overfitting — usually better inside an ensemble (Random Forest / Gradient Boosting)."
    ),
    ("gradient boosting", "xgboost", "boosting"): (
        "**Gradient Boosting** builds trees sequentially, each one correcting the errors of "
        "the previous ones. Often the top performer on tabular data (XGBoost/LightGBM are "
        "popular variants). Pros: high accuracy. Cons: more hyperparameters, slower to train."
    ),
    ("svm", "svr", "svc", "support vector"): (
        "**Support Vector Machines** find the boundary that maximizes the margin between "
        "classes (SVC) or fit within an error tube (SVR). Kernels enable non-linear "
        "boundaries. Pros: effective in high dimensions. Cons: slow on large datasets, "
        "needs feature scaling."
    ),
    ("knn", "k-nearest", "nearest neighbor"): (
        "**K-Nearest Neighbors** predicts from the K closest training points — majority "
        "vote for classification, average for regression. Pros: simple, no training phase. "
        "Cons: slow at prediction time, sensitive to scaling and irrelevant features."
    ),
    ("overfitting", "overfit"): (
        "**Overfitting** is when a model memorizes training data (including noise) and "
        "performs poorly on new data. Signs: high train accuracy, low test accuracy. "
        "Remedies: more data, simpler models, regularization, cross-validation, early stopping."
    ),
    ("which model", "best model", "what model", "choose a model", "recommend",
     "model should i use", "which algorithm", "what algorithm"): (
        "Quick guide: predicting a **number** → start with Linear Regression, then Random "
        "Forest / Gradient Boosting. Predicting a **category** → start with Logistic "
        "Regression, then Random Forest / Gradient Boosting. Small dataset with scaled "
        "features → SVM or KNN are worth a try. Upload your data and I can suggest a task "
        "type automatically from your target column."
    ),
    ("classification", "classify"): (
        "**Classification** predicts a discrete category (spam/not-spam, species, churn). "
        "Models here: Logistic Regression, Decision Tree, Random Forest, Gradient Boosting, "
        "SVC, and KNN. Metrics: accuracy, precision, recall, F1."
    ),
    ("regression", "predict a number", "continuous"): (
        "**Regression** predicts a continuous number (price, temperature, sales). Models "
        "here: Linear Regression, Decision Tree, Random Forest, Gradient Boosting, SVR, and "
        "KNN. Metrics: R², MAE, RMSE."
    ),
    ("preprocessing", "cleaning", "clean data", "preprocess"): (
        "**Preprocessing** in this app: fix messy headers, convert placeholder values and "
        "numeric-looking text, drop empty/constant columns, remove duplicates, impute or "
        "drop missing values, clip outliers, one-hot encode categoricals, and optionally "
        "scale numeric features. Head to the **Data & Preprocessing** tab to run it."
    ),
    ("missing value", "missing values", "impute", "imputation", "nan"): (
        "**Missing values** can be dropped (safe if few) or imputed — mean/median for "
        "numeric, mode for categorical. This app supports both strategies in the "
        "preprocessing step."
    ),
    ("scaling", "standardize", "normalize", "standardscaler", "min-max"): (
        "**Feature scaling** puts features on comparable ranges. Standard scaling (mean 0, "
        "std 1) or Min-Max (0-1) matters a lot for SVM and KNN; tree-based models don't "
        "need it. Choose a scaler in the preprocessing options."
    ),
    ("outlier", "outliers", "iqr"): (
        "**Outliers** are extreme values that can distort model fits, especially for linear "
        "models and KNN. This app can clip values beyond 1.5×IQR (the boxplot whisker rule) "
        "— enable it in the preprocessing options."
    ),
    ("metric", "metrics", "accuracy", "precision", "recall", "f1"): (
        "**Classification metrics**: accuracy = overall correct rate; precision = of "
        "predicted positives, how many were right; recall = of actual positives, how many "
        "were found; F1 = harmonic mean of precision and recall (good for imbalanced data)."
    ),
    ("r2", "r²", "rmse", "mae"): (
        "**Regression metrics**: R² = share of variance explained (1.0 is perfect); "
        "MAE = average absolute error in target units; RMSE = like MAE but penalizes "
        "large errors more."
    ),
    ("train test split", "test size", "split"): (
        "**Train/test split** holds out part of the data (default 20% here) to evaluate the "
        "model on unseen examples. Adjust it with the slider in the Model Training tab."
    ),
    ("feature importance",): (
        "**Feature importance** shows which columns most influence predictions. Tree models "
        "expose `feature_importances_`; linear models use coefficient magnitudes. This app "
        "plots the top 15 automatically after training."
    ),
    ("cross validation", "cross-validation", "k-fold"): (
        "**Cross-validation** splits data into K folds, trains K times each holding out one "
        "fold, and averages the scores — a more reliable estimate than a single split, at "
        "the cost of K× training time."
    ),
}

DATA_LOADED_PROMPT = (
    "\n\n📊 I can see you've uploaded **{name}**. Do you want me to analyse this data? "
    "Head to the **Data & Preprocessing** tab to clean it, then train models in "
    "**Model Training**."
)

FALLBACK = (
    "I can help with questions about the models in this app — linear/logistic regression, "
    "decision trees, random forests, gradient boosting, SVM, KNN — plus preprocessing, "
    "metrics, overfitting, and train/test splits.\n\n"
    "Or just get started: **upload your data in the Data & Preprocessing tab** and I'll "
    "help you analyse it. What do you want to analyse today?"
)

GREETINGS = ("hi", "hello", "hey", "good morning", "good afternoon", "good evening", "yo", "hola")


def _keyword_matches(msg: str, keyword: str) -> bool:
    """Whole-word match, so 'nan' does not fire inside 'financial'."""
    return re.search(r"(?<!\w)" + re.escape(keyword) + r"(?!\w)", msg) is not None


def _rule_based_reply(message: str) -> str:
    msg = message.lower().strip()

    if any(msg == g or msg.startswith(g + " ") or msg.startswith(g + "!") for g in GREETINGS):
        return (
            "Hello! I'm here to help you analyse your data with machine learning. "
            "**What do you want to analyse today?** You can ask about any ML model, "
            "or upload a dataset (CSV, Excel, JSON, Parquet) in the "
            "**Data & Preprocessing** tab."
        )

    # Specific knowledge first, generic workflow guidance after
    for keywords, answer in KNOWLEDGE_BASE.items():
        if any(_keyword_matches(msg, kw) for kw in keywords):
            return answer

    if any(_keyword_matches(msg, kw) for kw in
           ("analyse", "analyze", "analysis", "get started", "start", "upload")):
        return (
            "Great — here's the workflow:\n\n"
            "1. **Data & Preprocessing tab** — upload your file, review the profile, and "
            "run cleaning (missing values, duplicates, outliers, encoding, scaling).\n"
            "2. **Model Training tab** — pick a target column, a task "
            "(regression/classification), and a model, then train.\n"
            "3. Review metrics and charts, open the **code dropdown** to see the exact "
            "Python used, and **download the report** (PDF/HTML) and cleaned data.\n\n"
            "Do you want me to analyse your data? Upload it and I'll take it from there!"
        )

    return FALLBACK


# ---------------------------------------------------------------- entry ----


def respond(message: str, history: list | None = None,
            uploaded_file_name: str | None = None) -> str:
    """Return the assistant's reply: LLM if available, rules otherwise."""
    if os.environ.get("HF_TOKEN"):
        try:
            return _llm_reply(message, history, uploaded_file_name)
        except Exception:
            pass  # quota exhausted, provider down, etc. — fall back to rules

    reply = _rule_based_reply(message)
    if uploaded_file_name:
        reply += DATA_LOADED_PROMPT.format(name=uploaded_file_name)
    return reply