import gradio as gr
import os
from huggingface_hub import InferenceClient
from sentence_transformers import SentenceTransformer
import torch


# ============================================================
# LOAD KNOWLEDGE BASE
# ============================================================

with open("knowledge.txt", "r", encoding="utf-8") as f:
    knowledge_text = f.read()

# Split knowledge base into chunks
chunks = [
    chunk.strip()
    for chunk in knowledge_text.split("\n\n")
    if chunk.strip()
]


# ============================================================
# LOAD EMBEDDING MODEL
# ============================================================

embedder = SentenceTransformer("all-MiniLM-L6-v2")

# Precompute embeddings for knowledge-base chunks
chunk_embeddings = embedder.encode(
    chunks,
    convert_to_tensor=True
)


def get_relevant_context(query, top_k=3):
    """
    Find the most relevant pieces of the knowledge base
    for the user's question.
    """

    # Make sure top_k isn't larger than the number of chunks
    top_k = min(top_k, len(chunks))

    # Embed the user's question
    query_embedding = embedder.encode(
        query,
        convert_to_tensor=True
    )

    # Normalize query embedding
    query_embedding = query_embedding / query_embedding.norm()

    # Normalize knowledge-base embeddings
    norm_chunk_embeddings = (
        chunk_embeddings
        / chunk_embeddings.norm(dim=1, keepdim=True)
    )

    # Calculate cosine similarity
    similarities = torch.matmul(
        norm_chunk_embeddings,
        query_embedding
    )

    # Get the most relevant chunks
    top_k_indices = (
        torch.topk(similarities, k=top_k)
        .indices
        .cpu()
        .numpy()
    )

    # Combine the relevant chunks
    context = "\n\n".join(
        [chunks[i] for i in top_k_indices]
    )

    return context


# ============================================================
# GRADIO THEME
# ============================================================

custom_theme = gr.themes.Soft(
    primary_hue="green",
    secondary_hue="stone",
    neutral_hue="gray",
    spacing_size="md",
    radius_size="md",
    text_size="md",
    font=["Roboto", "sans-serif"],
    font_mono=["Roboto Mono", "monospace"],
)


# ============================================================
# HUGGING FACE CLIENT
# ============================================================

client = InferenceClient(
    token=os.environ["HF_TOKEN"],
    provider="auto"
)


# ============================================================
# CHATBOT FUNCTION
# ============================================================

def respond(message, history):

    # --------------------------------------------------------
    # SYSTEM PROMPT
    # --------------------------------------------------------

    system_prompt = """
You are ChaChingas, an AI financial literacy advisor for students
and low-income families.

Only answer questions about:
- budgeting
- saving
- debt
- credit
- credit card fraud
- investing
- banking
- personal finance

If a user asks about unrelated topics such as recipes, sports,
entertainment, homework, or other unrelated subjects, politely say:

"I'm here to help with money and budgeting—ask me anything about that!"

Use clear, simple language and keep answers relatively short.

When explaining budgeting, use the 50/30/20 rule:
- 50% for needs
- 30% for wants
- 20% for savings or debt

Be supportive, practical, and easy to understand.

Do not provide tax or legal advice.

Never ask for or handle sensitive personal financial information.

Use the provided knowledge base when it is relevant to the user's question.
If the knowledge base does not contain the answer, use your general
financial-literacy knowledge while staying within your restrictions.
"""

    # --------------------------------------------------------
    # GET RELEVANT KNOWLEDGE
    # --------------------------------------------------------

    context = get_relevant_context(
        message,
        top_k=3
    )

    # --------------------------------------------------------
    # BUILD MESSAGES
    # --------------------------------------------------------

    messages = [
        {
            "role": "system",
            "content": system_prompt
        }
    ]

    # Add relevant knowledge to the current question
    messages.append(
        {
            "role": "system",
            "content": (
                "Here is relevant information from the ChaChingas "
                "knowledge base:\n\n"
                + context
            )
        }
    )

    # --------------------------------------------------------
    # ADD CHAT HISTORY
    # --------------------------------------------------------

    if history:

        # Handle newer Gradio message format
        if isinstance(history[0], dict):

            for msg in history:
                if (
                    isinstance(msg, dict)
                    and msg.get("role") in ["user", "assistant"]
                    and msg.get("content")
                ):
                    messages.append(
                        {
                            "role": msg["role"],
                            "content": msg["content"]
                        }
                    )

        # Handle older Gradio tuple format
        else:

            for user_msg, assistant_msg in history:

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

                if assistant_msg:
                    messages.append(
                        {
                            "role": "assistant",
                            "content": assistant_msg
                        }
                    )

    # Add current question
    messages.append(
        {
            "role": "user",
            "content": message
        }
    )

    # --------------------------------------------------------
    # CALL HUGGING FACE
    # --------------------------------------------------------

    response = ""

    try:

        stream = client.chat_completion(
            model="deepseek-ai/DeepSeek-V3-0324",
            messages=messages,
            max_tokens=512,
            temperature=0.1,
            stream=True
        )

        # ----------------------------------------------------
        # PROCESS STREAM
        # ----------------------------------------------------

        for chunk in stream:

            # Some streaming chunks can have no choices
            if not chunk.choices:
                continue

            # Get content from the chunk
            token = chunk.choices[0].delta.content

            # Some chunks can have None as their content
            if not token:
                continue

            response += token

            # Send updated response to Gradio
            yield response

    except Exception as e:

        print("HUGGING FACE ERROR:", repr(e))

        yield (
            "Sorry, I ran into a problem connecting to the AI model. "
            "Please try again in a moment."
        )


# ============================================================
# GRADIO INTERFACE
# ============================================================

with gr.Blocks(theme=custom_theme) as demo:

    chatbot_interface = gr.ChatInterface(
        respond,
        examples=["Build a Budgeting Plan","Teach Me About Stocks","How Do I Set Up a Bank Account?", "How Do I Prevent Credit Card Fraud?"],
        title="ChaChingas",
        description="This is a financial literacy chatbot"
    )


# ============================================================
# LAUNCH
# ============================================================

demo.launch()