backend/app/ai/ollama_integration.py

import os
import requests
import json
import time

# ─── Backend credentials (read once at module load) ───────────────────────────
OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY", "")
OPENAI_MODEL   = os.environ.get("OPENAI_MODEL", "gpt-4o-mini")
GROQ_API_KEY   = os.environ.get("GROQ_API_KEY", "")
USE_GROQ       = bool(GROQ_API_KEY)

OLLAMA_URL = "http://127.0.0.1:11434"

# Check active backends once at load time to prevent timeout delays during requests.
# Priority: OpenAI → Groq → local Ollama
AI_BACKEND_AVAILABLE = False
if OPENAI_API_KEY or GROQ_API_KEY:
    AI_BACKEND_AVAILABLE = True
else:
    try:
        # Fast 0.5s ping to local Ollama
        response = requests.get(f"{OLLAMA_URL}/", timeout=0.5)
        AI_BACKEND_AVAILABLE = (response.status_code == 200)
    except Exception:
        AI_BACKEND_AVAILABLE = False

def has_active_ai_backend() -> bool:
    """Returns True if OpenAI, Groq, or local Ollama is active and reachable."""
    return AI_BACKEND_AVAILABLE

BANKING_KEYWORDS = [
    "account", "loan", "card", "balance",
    "transfer", "bank", "interest", "emi",
    "credit", "debit", "kyc", "upi", "cheque",
    "deposit", "fd", "rd", "branch", "ifsc",
    "transaction", "payment", "savings", "checking",
    "mortgage", "investment", "fintech", "wallet",
    "rate", "rates", "support", "customer", "care",
    "help", "contact", "helpline", "number", "call",
    "document", "required", "identity", "proof", "open"
]

SYSTEM_PROMPT = """You are BankBot, a professional banking assistant for Central Bank.
You ONLY answer banking-related questions. If the question is not related to banking, politely refuse.
Never answer questions about politics, sports, entertainment, coding, or personal advice.

CORE GUIDELINES:
1. ALWAYS communicate in {language}.
2. CONTEXT AWARENESS: Use the provided chat history to understand follow-up questions. For example, if the user asks "What is the interest rate?" and then "for home loan", you must understand they are asking about home loan interest rates.
3. CLARIFYING QUESTIONS: If a user's query is ambiguous (e.g., "how much?"), ask for missing details (e.g., "How much for what service? Balance check or loan EMI?").
4. CALCULATIONS: Perform financial calculations (EMI, Interest, Eligibility) if information is provided.
5. DOCUMENT ANALYSIS: If text from a PDF statement is provided, summarize it or answer specific questions about it.
6. PROFESSIONALISM: Maintain a helpful, formal, and secure tone."""

OLLAMA_URL = "http://127.0.0.1:11434"
DEFAULT_OLLAMA_MODEL = os.environ.get("OLLAMA_MODEL", "llama3:latest")


def is_banking_query(user_input):
    input_lower = user_input.lower()
    return any(word in input_lower for word in BANKING_KEYWORDS)


def get_active_backend():
    """Returns the highest-priority available backend name."""
    if OPENAI_API_KEY:
        return "openai"
    if USE_GROQ:
        return "groq"
    return "ollama"


def _build_messages(prompt, history=None, language="English"):
    sys_prompt = SYSTEM_PROMPT.format(language=language)
    messages = [{"role": "system", "content": sys_prompt}]

    if history:
        for msg in history[-10:]:
            if msg.get("role") and msg.get("content"):
                messages.append({"role": msg["role"], "content": msg["content"]})

    messages.append({"role": "user", "content": prompt})
    return messages


def _get_available_ollama_models():
    try:
        response = requests.get(f"{OLLAMA_URL}/api/tags", timeout=5)
        response.raise_for_status()
        data = response.json()
        return [model.get("name", "") for model in data.get("models", []) if model.get("name")]
    except Exception as e:
        print(f"Ollama model discovery error: {e}")
        return []


def _resolve_ollama_model(requested_model):
    available_models = _get_available_ollama_models()
    if not available_models:
        return requested_model

    if requested_model in available_models:
        return requested_model

    base_requested_model = requested_model.split(":", 1)[0]
    for candidate in available_models:
        if candidate.split(":", 1)[0] == base_requested_model:
            return candidate

    return available_models[0]


def _ollama_error_message(model, error):
    return (
        f"Ollama request failed for model '{model}': {error}. "
        "The Ollama server is reachable, but the model backend crashed internally. "
        "Try `ollama run llama3`, and if that fails restart Ollama with "
        "`taskkill /F /IM ollama.exe` followed by `ollama serve`."
    )


# ─── OpenAI Functions ────────────────────────────────────────────────────────

def get_openai_response(prompt, history=None, model=None, language="English"):
    """Fetches a response from the OpenAI API (gpt-4o-mini by default)."""
    if not OPENAI_API_KEY:
        return None
    try:
        from openai import OpenAI
        client = OpenAI(api_key=OPENAI_API_KEY)
        target_model = model or OPENAI_MODEL

        sys_prompt = SYSTEM_PROMPT.format(language=language)
        messages = [{"role": "system", "content": sys_prompt}]

        if history:
            for msg in history[-10:]:
                if msg.get("role") and msg.get("content"):
                    messages.append({"role": msg["role"], "content": msg["content"]})

        messages.append({"role": "user", "content": prompt})

        response = client.chat.completions.create(
            model=target_model,
            messages=messages,
            temperature=0.1,
            max_tokens=1000,
        )
        return response.choices[0].message.content
    except Exception as e:
        print(f"OpenAI Error: {e}")
        return None


def stream_openai_response(prompt, history=None, model=None, language="English"):
    """Yields streamed response chunks from the OpenAI API."""
    if not OPENAI_API_KEY:
        return
    try:
        from openai import OpenAI
        client = OpenAI(api_key=OPENAI_API_KEY)
        target_model = model or OPENAI_MODEL

        sys_prompt = SYSTEM_PROMPT.format(language=language)
        messages = [{"role": "system", "content": sys_prompt}]

        if history:
            for msg in history[-10:]:
                if msg.get("role") and msg.get("content"):
                    messages.append({"role": msg["role"], "content": msg["content"]})

        messages.append({"role": "user", "content": prompt})

        stream = client.chat.completions.create(
            model=target_model,
            messages=messages,
            temperature=0.1,
            max_tokens=1000,
            stream=True,
        )
        for chunk in stream:
            content = chunk.choices[0].delta.content
            if content:
                yield content
    except Exception as e:
        print(f"OpenAI Stream Error: {e}")


# ─── Groq AI Functions ────────────────────────────────────────────────────────

def get_groq_response(prompt, history=None, model="llama-3.3-70b-versatile", language="English"):
    """Fetches a response from Groq AI API."""
    try:
        from groq import Groq
        client = Groq(api_key=GROQ_API_KEY)

        sys_prompt = SYSTEM_PROMPT.format(language=language)
        messages = [{"role": "system", "content": sys_prompt}]

        if history:
            for msg in history[-10:]:
                if msg.get("role") and msg.get("content"):
                    messages.append({"role": msg["role"], "content": msg["content"]})

        messages.append({"role": "user", "content": prompt})

        response = client.chat.completions.create(
            model=model,
            messages=messages,
            temperature=0.1,
            max_tokens=1000,
        )
        return response.choices[0].message.content
    except Exception as e:
        print(f"Groq Error: {e}")
        return None


def stream_groq_response(prompt, history=None, model="llama-3.3-70b-versatile", language="English"):
    """Yields streamed response chunks from Groq AI API."""
    try:
        from groq import Groq
        client = Groq(api_key=GROQ_API_KEY)

        sys_prompt = SYSTEM_PROMPT.format(language=language)
        messages = [{"role": "system", "content": sys_prompt}]

        if history:
            for msg in history[-10:]:
                if msg.get("role") and msg.get("content"):
                    messages.append({"role": msg["role"], "content": msg["content"]})

        messages.append({"role": "user", "content": prompt})

        stream = client.chat.completions.create(
            model=model,
            messages=messages,
            temperature=0.1,
            max_tokens=1000,
            stream=True,
        )
        for chunk in stream:
            content = chunk.choices[0].delta.content
            if content:
                yield content
    except Exception as e:
        print(f"Groq Stream Error: {e}")
        yield None


# ─── Ollama Functions ─────────────────────────────────────────────────────────

def get_ollama_response(prompt, history=None, model=DEFAULT_OLLAMA_MODEL, language="English"):
    """Fetches a response from a local Ollama instance."""
    url = f"{OLLAMA_URL}/api/chat"
    resolved_model = _resolve_ollama_model(model)
    messages = _build_messages(prompt, history=history, language=language)

    payload = {
        "model": resolved_model,
        "messages": messages,
        "stream": False,
        "options": {"temperature": 0.1, "top_p": 0.9, "num_predict": 500}
    }

    try:
        # (connect_timeout, read_timeout) — cap total generation at 25s
        response = requests.post(url, json=payload, timeout=(5, 25))
        response.raise_for_status()
        data = response.json()
        return data.get("message", {}).get("content", "")
    except requests.exceptions.Timeout:
        # Don't retry on timeout — let the caller fall back to the next backend
        print(f"Ollama timed out for model '{resolved_model}'. Falling back to next backend.")
        return None
    except Exception as e:
        print(_ollama_error_message(resolved_model, e))
        if resolved_model != "llama3":
            return get_ollama_response(prompt, history, model="llama3", language=language)
        return None


def stream_ollama_response(prompt, history=None, model=DEFAULT_OLLAMA_MODEL, language="English"):
    """Yields chunks of the response from a local Ollama instance for streaming."""
    url = f"{OLLAMA_URL}/api/chat"
    resolved_model = _resolve_ollama_model(model)
    messages = _build_messages(prompt, history=history, language=language)

    payload = {
        "model": resolved_model,
        "messages": messages,
        "stream": True,
        "options": {"temperature": 0.1, "top_p": 0.9, "num_predict": 500}
    }

    try:
        # (connect_timeout, read_timeout) — cap total generation at 25s
        response = requests.post(url, json=payload, timeout=(5, 25), stream=True)
        response.raise_for_status()

        for line in response.iter_lines():
            if line:
                chunk = json.loads(line)
                if 'message' in chunk and 'content' in chunk['message']:
                    yield chunk['message']['content']
                if chunk.get('done'):
                    break
    except requests.exceptions.Timeout:
        # Don't retry on timeout — let the caller fall back to the next backend
        print(f"Ollama stream timed out for model '{resolved_model}'. Falling back to next backend.")
        return
    except Exception as e:
        print(_ollama_error_message(resolved_model, e))
        if resolved_model != "llama3":
            yield from stream_ollama_response(prompt, history, model="llama3", language=language)
        else:
            yield None


# ─── Unified Wrapper Functions ────────────────────────────────────────────────

def get_ai_response(prompt, history=None, language="English"):
    """
    Auto-selects the best available backend.
    Priority: OpenAI → Groq → Ollama
    Returns None only when all backends are unavailable.
    """
    if OPENAI_API_KEY:
        result = get_openai_response(prompt, history, language=language)
        if result:
            return result

    if USE_GROQ:
        result = get_groq_response(prompt, history, language=language)
        if result:
            return result

    return get_ollama_response(prompt, history, language=language)


def stream_ai_response(prompt, history=None, language="English"):
    """
    Auto-selects streaming from the best available backend.
    Priority: OpenAI → Groq → Ollama
    """
    if OPENAI_API_KEY:
        chunks = list(stream_openai_response(prompt, history, language=language))
        if chunks:
            yield from chunks
            return

    if USE_GROQ:
        chunks = list(stream_groq_response(prompt, history, language=language))
        if chunks:
            yield from chunks
            return

    yield from stream_ollama_response(prompt, history, language=language)


def check_ollama_connection():
    """Checks if the local Ollama server is running."""
    if USE_GROQ:
        return True
    try:
        response = requests.get(f"{OLLAMA_URL}/", timeout=2)
        return response.status_code == 200
    except:
        return False