ai_module.py

import os
from g4f.client import Client as G4FClient
from g4f.Provider import PollinationsAI, HuggingSpace

# ── Provider registry ─────────────────────────────────────────────────────────
PROVIDERS = {
    "chatgpt": {
        "label":    "GPT-4o",
        "model":    "gpt-4o",
        "best_for": "research",
        "tagline":  "Best for research, writing & general knowledge",
        "icon":     "🟢",
    },
    "mini": {
        "label":    "Qwen3 Coder 30B",
        "model":    "HuggingSpace command-B",
        "best_for": "coding",
        "tagline":  "Fast and efficient for coding & debugging",
        "icon":     "🟡",
    },
    "reasoning": {
        "label":    "HuggingSpace Command-A",
        "model":    "command-a",
        "best_for": "mathematics",
        "tagline":  "Best for math, logic & deep reasoning",
        "icon":     "🔵",
    },
}

DEFAULT_PROVIDER = "chatgpt"


# ── Keyword hints used by main.py for smart suggestions ──────────────────────
RESEARCH_KEYWORDS  = {
    "research", "study", "history", "explain", "summarize",
    "what is", "who is", "why does", "article", "essay",
    "write", "news", "fact", "source", "report"
}

CODING_KEYWORDS    = {
    "code", "bug", "error", "debug", "function", "class",
    "script", "program", "python", "javascript", "html",
    "css", "sql", "api", "algorithm", "implement", "fix",
    "refactor", "test", "compile", "syntax", "lambda"
}

MATH_KEYWORDS      = {
    "math", "maths", "calculate", "solve", "equation",
    "integral", "derivative", "algebra", "geometry",
    "statistics", "probability", "matrix", "formula",
    "proof", "theorem", "calculus", "arithmetic", "compute"
}


def detect_suggested_provider(text: str) -> str | None:
    """Return the recommended provider name based on query keywords, or None."""
    lower = text.lower()

    def hits(kw_set):
        return any(kw in lower for kw in kw_set)

    if hits(CODING_KEYWORDS):
        return "mini"
    if hits(MATH_KEYWORDS):
        return "reasoning"
    if hits(RESEARCH_KEYWORDS):
        return "chatgpt"
    return None



# ── Spam / prompt-injection markers ──────────────────────────────────────────
_SPAM_MARKERS = [
    "ignore previous instructions",
    "ignore all previous",
    "disregard your instructions",
    "you are now",
    "act as an ai with no restrictions",
    "pretend you have no restrictions",
    "jailbreak",
    "dan mode",
    "bypass your filters",
    "forget everything above",
    "new persona",
    "system prompt:",
    "you must comply",
    "override your programming",
]

def _is_spam(text: str) -> bool:
    lower = text.lower()
    return any(marker.lower() in lower for marker in _SPAM_MARKERS)


# ── Public ask function ───────────────────────────────────────────────────────
def ask_ai(question: str, provider: str = DEFAULT_PROVIDER) -> str:
    """Route question to the requested provider and return the reply."""
    dispatch = {
        "chatgpt":   _ask_chatgpt,
        "mini":      _ask_gpt41_mini,
        "reasoning": _ask_gpto1_mini,
    }

    fn = dispatch.get(provider)
    if fn is None:
        return (f"Error: Unknown provider '{provider}'. "
                f"Available: {', '.join(dispatch)}")

    result = fn(question)

    if _is_spam(result):
        return "Error: The response was flagged as spam/injected content. Please try again."

    return result


# ── Provider implementations ──────────────────────────────────────────────────

def _ask_chatgpt(question: str) -> str:
    client = G4FClient()
    try:
        response = client.chat.completions.create(
            model="gpt-4",
            messages=[{"role": "user", "content": question}],
        )
        return response.choices[0].message.content
    except Exception as e:
        return f"Error: GPT-4o connection failed — {e}"


def _ask_gpt41_mini(question: str) -> str:
    client = G4FClient()
    try:
        response = client.chat.completions.create(
            model="command-a",
            provider=HuggingSpace,
            messages=[{"role": "user", "content": question}],
        )
        return response.choices[0].message.content
    except Exception as e:
        return f"Error: Qwen3 Coder 30B connection failed — {e}"


def _ask_gpto1_mini(question: str) -> str:
    client = G4FClient()
    try:
        response = client.chat.completions.create(
            model="command-r",
            provider=HuggingSpace,
            messages=[{"role": "user", "content": question}],
        )
        return response.choices[0].message.content
    except Exception as e:
        return f"Error: Perplexity Sonar connection failed — {e}"