backend/app/ai/chat.py

import json
import os
from threading import Lock
from sqlalchemy.orm import Session
from app.database.models import User, Account, Transaction, Goal, Investment, Subscription
from app.ai.behavior import analyze_spending_behavior
from app.ai.coaching import calculate_financial_health_score
from app.ai.ollama_integration import get_groq_response, get_ollama_response, stream_groq_response, stream_ollama_response

# Thread-safe chatbot memory storage
class ChatMemoryManager:
    def __init__(self):
        self._history = {}
        self._lock = Lock()

    def get_history(self, user_id: str):
        with self._lock:
            if user_id not in self._history:
                self._history[user_id] = []
            return self._history[user_id]

    def add_message(self, user_id: str, role: str, content: str):
        with self._lock:
            if user_id not in self._history:
                self._history[user_id] = []
            self._history[user_id].append({"role": role, "content": content})
            # Limit history to last 12 messages (6 rounds)
            if len(self._history[user_id]) > 12:
                self._history[user_id] = self._history[user_id][-12:]

    def clear_history(self, user_id: str):
        with self._lock:
            if user_id in self._history:
                self._history[user_id] = []

chat_memory = ChatMemoryManager()

def build_user_context_string(db: Session, user_id: str) -> str:
    """
    Queries database for a user's entire financial situation to construct a precise system context.
    """
    user = db.query(User).filter(User.id == user_id).first()
    if not user:
        return "No user information available."
        
    accounts = db.query(Account).filter(Account.user_id == user_id).all()
    total_balance = sum(acc.balance for acc in accounts)
    account_details = [f"{acc.type.capitalize()} Account: ${acc.balance:,.2f}" for acc in accounts]
    
    goals = db.query(Goal).filter(Goal.user_id == user_id).all()
    goals_details = [f"Goal '{g.title}': Target ${g.target_amount:,.2f}, Saved ${g.current_amount:,.2f}" for g in goals]
    
    investments = db.query(Investment).filter(Investment.user_id == user_id).all()
    investments_details = [f"{i.asset_name} ({i.type}): invested ${i.amount_invested:,.2f}, Current Value ${i.current_value:,.2f}" for i in investments]
    
    subs = db.query(Subscription).filter(Subscription.user_id == user_id, Subscription.active == True).all()
    subs_details = [f"{s.merchant}: ${s.amount:,.2f}/{s.billing_cycle}" for s in subs]
    
    # Run behavioral diagnostics
    behavior = analyze_spending_behavior(db, user_id)
    behavior_insights = behavior.get("insights", [])
    
    # Financial Score
    score_data = calculate_financial_health_score(db, user_id)
    financial_score = score_data.get("overall_score", 50)
    
    context = f"""
    User Profile:
    - Name: {user.profile_data.get('name', 'Client')}
    - Financial Personality: {user.financial_personality}
    - Financial Health Score: {financial_score:.0f}/100
    
    Balances:
    {', '.join(account_details) if account_details else 'No active bank accounts'}
    - Total Liquid Capital: ${total_balance:,.2f}
    
    Financial Goals:
    {'; '.join(goals_details) if goals_details else 'None established'}
    
    Active Portfolio:
    {'; '.join(investments_details) if investments_details else 'No active investments'}
    
    Active Subscriptions:
    {'; '.join(subs_details) if subs_details else 'No active subscriptions'}
    
    Diagnostics & Behavior:
    - {'; '.join(behavior_insights)}
    - Late night spending occurrences: {behavior.get('metrics', {}).get('late_night_count', 0)}
    - Weekend spending ratio: {behavior.get('metrics', {}).get('weekend_pct', 0.0)}%
    """
    return context

def get_contextual_system_prompt(db: Session, user_id: str) -> str:
    """
    Constructs a highly specific system prompt containing the user's financial profile.
    """
    financial_context = build_user_context_string(db, user_id)
    
    system_prompt = f"""You are BankBot, an elite AI Financial Analyst, Wealth Advisor, and Predictive Banking Engine.
    You communicate with the user, providing highly personalized, concise, and mathematically rigorous answers.
    You have direct, read-only access to the client's current financial profile and database records.

    CURRENT USER PORTFOLIO DATA:
    {financial_context}

    CORE PRINCIPLES:
    1. NEVER behave like a generic chatbot. Avoid generic suggestions like "save more money". Use real numbers, calculate percentages, and suggest specific actions based on the client's data.
    2. Respond with the authority and brevity of a Bloomberg Terminal analyst.
    3. Keep your answers brief, actionable, and financially meaningful (typically 2-4 sentences max).
    4. If the user asks a question about their spending, goals, or predictions, use the portfolio data above.
    5. Always remain helpful, professional, and secure.
    """
    return system_prompt

def get_offline_chat_fallback(db: Session, user_id: str, prompt: str) -> str:
    """
    Generates a localized, rule-grounded financial analyst reply when AI engines are offline.
    """
    user = db.query(User).filter(User.id == user_id).first()
    persona = user.financial_personality if user else "Saver"
    
    prompt_lower = prompt.lower()
    
    if "discipline" in prompt_lower or "spend" in prompt_lower or "budget" in prompt_lower:
        score_data = calculate_financial_health_score(db, user_id)
        discipline_score = score_data.get("categories", {}).get("spending_discipline", {}).get("score", 10)
        return (
            f"As a {persona}, your spending discipline score stands at {discipline_score:.0f}/20. "
            f"Analysis of your transaction history shows discretionary spikes. "
            "To optimize your cashflow surplus, establish a strict 20% savings buffer prior to discretionary outflow."
        )
    elif "investment" in prompt_lower or "portfolio" in prompt_lower or "grow" in prompt_lower:
        investments = db.query(Investment).filter(Investment.user_id == user_id).all()
        inv_total = sum(i.current_value for i in investments)
        return (
            f"Your current investment portfolio valuation stands at ${inv_total:,.2f}. "
            "Based on asset performance, shifting 15% of your net checking surplus into stock index funds "
            "will counter inflation and capture a projected 8% compound annual return."
        )
    else:
        score_data = calculate_financial_health_score(db, user_id)
        score = score_data.get("overall_score", 50)
        return (
            f"Wealth Advisor assessment: Your overall Financial Health Score is {score:.0f}/100. "
            "Liquidity is stable, but subscription and discretionary leakages are tempering compounding growth. "
            "Audit duplicate subscriptions and automate goal savings to enhance your trajectory."
        )

def get_chat_response(db: Session, user_id: str, prompt: str) -> str:
    """
    Returns an HTTP conversational response grounded in database context.
    """
    from app.ai.ollama_integration import has_active_ai_backend
    
    if not has_active_ai_backend():
        fallback_msg = get_offline_chat_fallback(db, user_id, prompt)
        chat_memory.add_message(user_id, "user", prompt)
        chat_memory.add_message(user_id, "assistant", fallback_msg)
        return fallback_msg
        
    sys_prompt = get_contextual_system_prompt(db, user_id)
    history = chat_memory.get_history(user_id)
    
    # Construct complete prompt for underlying backend
    full_messages = [{"role": "system", "content": sys_prompt}]
    for msg in history:
        full_messages.append({"role": msg["role"], "content": msg["content"]})
    full_messages.append({"role": "user", "content": prompt})
    
    # Determine backend
    OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY", "")
    GROQ_API_KEY = os.environ.get("GROQ_API_KEY", "")
    
    response_content = None
    
    if OPENAI_API_KEY:
        try:
            from openai import OpenAI
            client = OpenAI(api_key=OPENAI_API_KEY)
            res = client.chat.completions.create(
                model="gpt-4o-mini",
                messages=full_messages,
                temperature=0.1,
                max_tokens=500
            )
            response_content = res.choices[0].message.content
        except Exception as e:
            print(f"OpenAI error in chat: {e}")
            
    if not response_content and GROQ_API_KEY:
        try:
            response_content = get_groq_response(prompt, history=history, language="English")
        except Exception as e:
            print(f"Groq error in chat: {e}")
            
    if not response_content:
        # Fallback to local Ollama integration
        try:
            response_content = get_ollama_response(prompt, history=history, language="English")
        except Exception as e:
            print(f"Ollama error in chat: {e}")
            
    if not response_content:
        response_content = get_offline_chat_fallback(db, user_id, prompt)
        
    # Save conversation
    chat_memory.add_message(user_id, "user", prompt)
    chat_memory.add_message(user_id, "assistant", response_content)
    
    return response_content

def stream_chat_response(db: Session, user_id: str, prompt: str):
    """
    Generates streaming chunks for WebSocket or HTTP SSE.
    """
    from app.ai.ollama_integration import has_active_ai_backend
    
    if not has_active_ai_backend():
        fallback_msg = get_offline_chat_fallback(db, user_id, prompt)
        chat_memory.add_message(user_id, "user", prompt)
        chat_memory.add_message(user_id, "assistant", fallback_msg)
        # Yield words slowly to simulate streaming
        import time
        for word in fallback_msg.split(" "):
            yield word + " "
            time.sleep(0.05)
        return
        
    sys_prompt = get_contextual_system_prompt(db, user_id)
    history = chat_memory.get_history(user_id)
    
    full_messages = [{"role": "system", "content": sys_prompt}]
    for msg in history:
        full_messages.append({"role": msg["role"], "content": msg["content"]})
    full_messages.append({"role": "user", "content": prompt})
    
    # Save user message to history
    chat_memory.add_message(user_id, "user", prompt)
    
    OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY", "")
    GROQ_API_KEY = os.environ.get("GROQ_API_KEY", "")
    
    complete_reply = ""
    
    if OPENAI_API_KEY:
        try:
            from openai import OpenAI
            client = OpenAI(api_key=OPENAI_API_KEY)
            stream = client.chat.completions.create(
                model="gpt-4o-mini",
                messages=full_messages,
                temperature=0.1,
                max_tokens=500,
                stream=True
            )
            for chunk in stream:
                content = chunk.choices[0].delta.content
                if content:
                    complete_reply += content
                    yield content
            # Save assistant message once streaming completes
            chat_memory.add_message(user_id, "assistant", complete_reply)
            return
        except Exception as e:
            print(f"OpenAI streaming error: {e}")
            
    if GROQ_API_KEY:
        try:
            for chunk in stream_groq_response(prompt, history=history, language="English"):
                if chunk:
                    complete_reply += chunk
                    yield chunk
            chat_memory.add_message(user_id, "assistant", complete_reply)
            return
        except Exception as e:
            print(f"Groq streaming error: {e}")
            
    # Fallback to local Ollama stream
    try:
        for chunk in stream_ollama_response(prompt, history=history, language="English"):
            if chunk:
                complete_reply += chunk
                yield chunk
        chat_memory.add_message(user_id, "assistant", complete_reply)
    except Exception as e:
        print(f"Ollama streaming error: {e}")
        fallback_msg = get_offline_chat_fallback(db, user_id, prompt)
        yield fallback_msg
        chat_memory.add_message(user_id, "assistant", fallback_msg)