backend/app/ai/chat.py

import json
import os
from sqlalchemy.orm import Session
from datetime import datetime, timezone

from app.database.models import (
    User, Account, Transaction, Goal, Investment, Subscription,
    ChatMessage, ChatSession,
)
from app.ai.behavior import analyze_spending_behavior
from app.ai.coaching import calculate_financial_health_score

CONTEXT_LIMIT = 12
MESSAGES_PER_SESSION_LIMIT = 200
SESSIONS_PER_USER_LIMIT = 50


def _title_from_prompt(text: str) -> str:
    cleaned = " ".join(text.strip().split())
    if len(cleaned) <= 48:
        return cleaned or "New chat"
    return cleaned[:48] + "…"


class ChatMemoryManager:
    """Session-scoped chat persistence with multi-conversation support."""

    def get_session(self, db: Session, user_id: str, session_id: str) -> ChatSession | None:
        return (
            db.query(ChatSession)
            .filter(ChatSession.id == session_id, ChatSession.user_id == user_id)
            .first()
        )

    def create_session(self, db: Session, user_id: str, title: str = "New chat") -> ChatSession:
        session = ChatSession(user_id=user_id, title=title or "New chat")
        db.add(session)
        db.commit()
        db.refresh(session)
        self._trim_old_sessions(db, user_id)
        return session

    def list_sessions(self, db: Session, user_id: str, limit: int = 30) -> list[dict]:
        limit = min(max(limit, 1), SESSIONS_PER_USER_LIMIT)
        sessions = (
            db.query(ChatSession)
            .filter(ChatSession.user_id == user_id)
            .order_by(ChatSession.updated_at.desc())
            .limit(limit)
            .all()
        )
        result = []
        for s in sessions:
            last = (
                db.query(ChatMessage)
                .filter(ChatMessage.session_id == s.id)
                .order_by(ChatMessage.created_at.desc())
                .first()
            )
            count = db.query(ChatMessage).filter(ChatMessage.session_id == s.id).count()
            result.append({
                "id": s.id,
                "title": s.title,
                "created_at": s.created_at.isoformat() if s.created_at else None,
                "updated_at": s.updated_at.isoformat() if s.updated_at else None,
                "message_count": count,
                "preview": (last.content[:80] + "…") if last and len(last.content) > 80 else (last.content if last else ""),
            })
        return result

    def delete_session(self, db: Session, user_id: str, session_id: str) -> bool:
        session = self.get_session(db, user_id, session_id)
        if not session:
            return False
        db.delete(session)
        db.commit()
        return True

    def ensure_session(self, db: Session, user_id: str, session_id: str | None) -> ChatSession:
        if session_id:
            session = self.get_session(db, user_id, session_id)
            if session:
                return session
        return self.create_session(db, user_id)

    def get_history(self, db: Session, session_id: str) -> list[dict]:
        rows = (
            db.query(ChatMessage)
            .filter(ChatMessage.session_id == session_id)
            .order_by(ChatMessage.created_at.desc())
            .limit(CONTEXT_LIMIT)
            .all()
        )
        return [{"role": r.role, "content": r.content} for r in reversed(rows)]

    def list_messages(self, db: Session, session_id: str, limit: int = 100) -> list[dict]:
        limit = min(max(limit, 1), MESSAGES_PER_SESSION_LIMIT)
        rows = (
            db.query(ChatMessage)
            .filter(ChatMessage.session_id == session_id)
            .order_by(ChatMessage.created_at.asc())
            .limit(limit)
            .all()
        )
        return [
            {
                "id": r.id,
                "role": r.role,
                "content": r.content,
                "created_at": r.created_at.isoformat() if r.created_at else None,
            }
            for r in rows
        ]

    def add_message(
        self, db: Session, session: ChatSession, user_id: str, role: str, content: str
    ) -> None:
        if not content or not content.strip():
            return
        text = content.strip()
        db.add(
            ChatMessage(
                user_id=user_id,
                session_id=session.id,
                role=role,
                content=text,
            )
        )
        session.updated_at = datetime.now(timezone.utc)
        if role == "user" and session.title in ("New chat", "Previous conversation"):
            session.title = _title_from_prompt(text)
        db.commit()
        self._trim_old_messages(db, session.id)

    def clear_session_messages(self, db: Session, user_id: str, session_id: str) -> bool:
        session = self.get_session(db, user_id, session_id)
        if not session:
            return False
        db.query(ChatMessage).filter(ChatMessage.session_id == session_id).delete()
        session.title = "New chat"
        session.updated_at = datetime.now(timezone.utc)
        db.commit()
        return True

    def _trim_old_messages(self, db: Session, session_id: str) -> None:
        total = db.query(ChatMessage).filter(ChatMessage.session_id == session_id).count()
        if total <= MESSAGES_PER_SESSION_LIMIT:
            return
        excess = total - MESSAGES_PER_SESSION_LIMIT
        oldest = (
            db.query(ChatMessage)
            .filter(ChatMessage.session_id == session_id)
            .order_by(ChatMessage.created_at.asc())
            .limit(excess)
            .all()
        )
        ids = [m.id for m in oldest]
        if ids:
            db.query(ChatMessage).filter(ChatMessage.id.in_(ids)).delete(synchronize_session=False)
            db.commit()

    def _trim_old_sessions(self, db: Session, user_id: str) -> None:
        total = db.query(ChatSession).filter(ChatSession.user_id == user_id).count()
        if total <= SESSIONS_PER_USER_LIMIT:
            return
        excess = total - SESSIONS_PER_USER_LIMIT
        oldest = (
            db.query(ChatSession)
            .filter(ChatSession.user_id == user_id)
            .order_by(ChatSession.updated_at.asc())
            .limit(excess)
            .all()
        )
        for s in oldest:
            db.delete(s)
        db.commit()


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} "
        f"({(g.current_amount/g.target_amount*100):.0f}% complete)"
        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}, "
        f"current value ${i.current_value:,.2f} "
        f"({'▲' if i.current_value >= i.amount_invested else '▼'}"
        f"{abs(i.current_value - i.amount_invested):,.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]
    monthly_sub_cost = sum(
        s.amount if s.billing_cycle == "monthly" else s.amount / 12
        for s in subs
    )

    # Recent transactions (last 10)
    from app.database.models import Account as Acct
    account_ids = [a.id for a in accounts]
    recent_txns = (
        db.query(Transaction)
        .filter(Transaction.account_id.in_(account_ids))
        .order_by(Transaction.timestamp.desc())
        .limit(10)
        .all()
    ) if account_ids else []
    txn_lines = [
        f"{t.merchant or 'Unknown'} ({t.category or 'Other'}): "
        f"{'+'if t.type=='credit' else '-'}${t.amount:,.2f}"
        for t in recent_txns
    ]

    # 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)
    score_categories = score_data.get("categories", {})
    score_lines = [
        f"{k.replace('_',' ').title()}: {v.get('score',0):.0f}/{v.get('max',20)}"
        for k, v in score_categories.items()
    ]

    context = f"""
USER PROFILE:
- Name: {user.profile_data.get('name', 'Client')}
- Email: {user.email}
- Financial Personality: {user.financial_personality}
- Financial Health Score: {financial_score:.0f}/100
- Score Breakdown: {', '.join(score_lines) if score_lines else 'N/A'}

ACCOUNT BALANCES:
{chr(10).join('  - ' + d for d in account_details) if account_details else '  - No active accounts'}
- Total Liquid Capital: ${total_balance:,.2f}

FINANCIAL GOALS:
{chr(10).join('  - ' + d for d in goals_details) if goals_details else '  - None established'}

INVESTMENT PORTFOLIO:
{chr(10).join('  - ' + d for d in investments_details) if investments_details else '  - No active investments'}
- Total Portfolio Value: ${sum(i.current_value for i in investments):,.2f}

ACTIVE SUBSCRIPTIONS (${monthly_sub_cost:,.2f}/month total):
{chr(10).join('  - ' + d for d in subs_details) if subs_details else '  - None'}

RECENT TRANSACTIONS (last 10):
{chr(10).join('  - ' + t for t in txn_lines) if txn_lines else '  - No recent transactions'}

BEHAVIORAL ANALYSIS:
{chr(10).join('  - ' + i for i in behavior_insights) if behavior_insights else '  - Insufficient data'}
- Late night spending occurrences: {behavior.get('metrics', {}).get('late_night_count', 0)}
- Weekend spending ratio: {behavior.get('metrics', {}).get('weekend_pct', 0.0):.1f}%
"""
    return context

def get_contextual_system_prompt(db: Session, user_id: str, language: str = "English") -> 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, a personal AI financial advisor with DIRECT ACCESS to this user's real bank account data.

CRITICAL RULES — NEVER BREAK THESE:
1. You ALREADY HAVE the user's complete financial data below. NEVER say "I don't have access to your account" or "I would need more information". You have everything.
2. ALWAYS answer using the EXACT numbers from the data below. Quote balances, amounts, percentages directly.
3. Be concise and specific — 2-5 sentences max. No generic advice.
4. If asked about balance, quote the exact figure. If asked about spending, reference actual transactions. If asked about goals, quote exact progress.
5. Speak like a personal banker who knows this client's finances intimately.
6. ALWAYS communicate in {language}.

THIS USER'S LIVE FINANCIAL DATA:
{financial_context}

You have read-only access to this data. Use it to answer every question with precision."""
    return system_prompt

def get_offline_chat_fallback(db: Session, user_id: str, prompt: str, language: str = "English") -> 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, session_id: str | None = None, language: str = "English"
) -> tuple[str, str]:
    """
    Returns (assistant reply, session_id) grounded in database context.
    """
    from app.ai.ollama_integration import has_active_ai_backend

    session = chat_memory.ensure_session(db, user_id, session_id)

    if not has_active_ai_backend():
        fallback_msg = get_offline_chat_fallback(db, user_id, prompt, language=language)
        chat_memory.add_message(db, session, user_id, "user", prompt)
        chat_memory.add_message(db, session, user_id, "assistant", fallback_msg)
        return fallback_msg, session.id

    # Build full message list: system prompt with real account data + history + new message
    sys_prompt = get_contextual_system_prompt(db, user_id, language=language)
    history = chat_memory.get_history(db, session.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})

    OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY", "")
    GROQ_API_KEY = os.environ.get("GROQ_API_KEY", "") or os.environ.get("GROQ_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=600,
            )
            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:
            # Pass full_messages directly so account context is included
            from groq import Groq
            client = Groq(api_key=GROQ_API_KEY)
            res = client.chat.completions.create(
                model="llama-3.3-70b-versatile",
                messages=full_messages,
                temperature=0.1,
                max_tokens=600,
            )
            response_content = res.choices[0].message.content
        except Exception as e:
            print(f"Groq error in chat: {e}")

    if not response_content:
        try:
            from app.ai.ollama_integration import get_ollama_response
            response_content = get_ollama_response(prompt, history=history, language=language)
        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, language=language)

    chat_memory.add_message(db, session, user_id, "user", prompt)
    chat_memory.add_message(db, session, user_id, "assistant", response_content)
    return response_content, session.id

def stream_chat_response(
    db: Session, user_id: str, prompt: str, session_id: str | None = None, language: str = "English"
):
    """
    Generates streaming chunks for WebSocket or HTTP SSE.
    """
    from app.ai.ollama_integration import has_active_ai_backend

    session = chat_memory.ensure_session(db, user_id, session_id)

    if not has_active_ai_backend():
        fallback_msg = get_offline_chat_fallback(db, user_id, prompt, language=language)
        chat_memory.add_message(db, session, user_id, "user", prompt)
        chat_memory.add_message(db, session, user_id, "assistant", fallback_msg)
        import time
        for word in fallback_msg.split(" "):
            yield word + " "
            time.sleep(0.05)
        return

    sys_prompt = get_contextual_system_prompt(db, user_id, language=language)
    history = chat_memory.get_history(db, session.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})

    chat_memory.add_message(db, session, user_id, "user", prompt)

    OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY", "")
    GROQ_API_KEY = os.environ.get("GROQ_API_KEY", "") or os.environ.get("GROQ_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=600,
                stream=True,
            )
            for chunk in stream:
                content = chunk.choices[0].delta.content
                if content:
                    complete_reply += content
                    yield content
            chat_memory.add_message(db, session, user_id, "assistant", complete_reply)
            return
        except Exception as e:
            print(f"OpenAI streaming error: {e}")

    if GROQ_API_KEY:
        try:
            from groq import Groq
            client = Groq(api_key=GROQ_API_KEY)
            stream = client.chat.completions.create(
                model="llama-3.3-70b-versatile",
                messages=full_messages,
                temperature=0.1,
                max_tokens=600,
                stream=True,
            )
            for chunk in stream:
                content = chunk.choices[0].delta.content
                if content:
                    complete_reply += content
                    yield content
            chat_memory.add_message(db, session, user_id, "assistant", complete_reply)
            return
        except Exception as e:
            print(f"Groq streaming error: {e}")

    try:
        from app.ai.ollama_integration import stream_ollama_response
        for chunk in stream_ollama_response(prompt, history=history, language=language):
            if chunk:
                complete_reply += chunk
                yield chunk
        chat_memory.add_message(db, session, 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, language=language)
        yield fallback_msg
        chat_memory.add_message(db, session, user_id, "assistant", fallback_msg)