import gradio as gr
import torch
import pdfplumber
import pandas as pd
import numpy as np
import re
import faiss
import json
import os
import numpy_financial as npf

from functools import lru_cache

from groq import Groq
from sentence_transformers import SentenceTransformer

# =========================================================
# PERFORMANCE SETTINGS
# =========================================================

os.environ["TOKENIZERS_PARALLELISM"] = "false"

torch.set_num_threads(2)

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# =========================================================
# GROQ CLIENT
# =========================================================

client = Groq(
    api_key=os.environ.get("GROQ_API_KEY")
)

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

embedder = SentenceTransformer(
    "BAAI/bge-small-en-v1.5",
    device=DEVICE
)

# =========================================================
# FAISS DATABASE
# =========================================================

FAISS_FOLDER = "finance_faiss"

faiss_index = None
doc_chunks = []

def load_faiss_database():

    global faiss_index
    global doc_chunks

    try:

        possible_index_paths = [
            "finance_faiss.index",
            os.path.join(FAISS_FOLDER, "finance_faiss.index")
        ]

        possible_chunk_paths = [
            "documents.json",
            os.path.join(FAISS_FOLDER, "documents.json")
        ]

        index_path = None
        chunks_path = None

        for path in possible_index_paths:

            if os.path.exists(path):

                index_path = path
                break

        for path in possible_chunk_paths:

            if os.path.exists(path):

                chunks_path = path
                break

        if index_path is None:

            print("⚠️ finance_faiss.index not found")
            return

        if chunks_path is None:

            print("⚠️ documents.json not found")
            return

        print(f"Loading FAISS index from: {index_path}")
        print(f"Loading documents from: {chunks_path}")

        faiss_index = faiss.read_index(index_path)

        with open(chunks_path, "r", encoding="utf-8") as f:

            doc_chunks = json.load(f)

        print("✅ FAISS database loaded successfully")
        print(f"✅ Total chunks loaded: {len(doc_chunks)}")

    except Exception as e:

        print("❌ Failed loading FAISS database")
        print(str(e))

load_faiss_database()

# =========================================================
# SYSTEM PROMPT
# =========================================================

SYSTEM_PROMPT = """
You are FinBot, an advanced conversational financial AI assistant.

Your behavior:
- Be conversational and intelligent
- Maintain conversation memory
- Understand follow-up questions
- Give detailed and complete answers
- Explain concepts clearly
- Use examples whenever useful
- Never hallucinate financial facts
- If unsure, say so honestly

Formatting:
- Use markdown formatting
- Use bullet points where useful
- Use numbered explanations when needed
- Keep answers readable and structured
"""

# =========================================================
# SETTINGS
# =========================================================

MAX_HISTORY_TURNS = 10
MAX_CONTEXT_CHARS = 2200

# =========================================================
# CLEAN RESPONSE
# =========================================================

def clean_response(text):

    text = re.sub(r"<\|.*?\|>", "", text)

    text = text.replace("assistant", "")

    text = text.strip()

    return text

# =========================================================
# FORMAT RESPONSE
# =========================================================

def format_response(text):

    sections = [
        "Revenue",
        "Profitability",
        "Assets",
        "Liabilities",
        "Equity",
        "Key Insights",
        "Formula",
        "Final Answer",
        "Financial Strengths",
        "Financial Weaknesses"
    ]

    for section in sections:

        text = text.replace(
            f"{section}:",
            f"\n### {section}\n"
        )

    return text.strip()

# =========================================================
# GENERATE RESPONSE
# =========================================================

def generate_response(messages, max_tokens=500):

    try:

        completion = client.chat.completions.create(

            model="yakul259/finance-rag-chatbot",

            messages=messages,

            temperature=0.4,

            top_p=0.9,

            max_tokens=max_tokens,

            presence_penalty=0.1,

            frequency_penalty=0.1,

            stop=[
                "<|eot_id|>",
                "User:"
            ]
        )

        response = completion.choices[0].message.content

        return clean_response(response)

    except Exception as e:

        return f"❌ Error generating response: {str(e)}"

# =========================================================
# FINANCIAL CALCULATOR
# =========================================================

def solve_financial_query(query):

    q = query.lower()

    try:

        # FUTURE VALUE

        if "future value" in q:

            principal = float(
                re.search(r'(\d+)', q).group(1)
            )

            rate = float(
                re.search(r'(\d+)%', q).group(1)
            ) / 100

            years = float(
                re.search(r'for (\d+)', q).group(1)
            )

            fv = npf.fv(
                rate=rate,
                nper=years,
                pmt=0,
                pv=-principal
            )

            return f"""
Formula:
FV = PV × (1 + r)^n

Values:
PV = {principal}
r = {rate * 100}%
n = {years}

Final Answer:
Future Value = {fv:.2f}
"""

        # PRESENT VALUE

        elif "present value" in q:

            values = re.findall(r'\d+', q)

            future_value = float(values[0])

            rate = float(values[1]) / 100

            years = float(values[2])

            pv = npf.pv(
                rate=rate,
                nper=years,
                pmt=0,
                fv=-future_value
            )

            return f"""
Formula:
PV = FV / (1 + r)^n

Values:
FV = {future_value}
r = {rate * 100}%
n = {years}

Final Answer:
Present Value = {abs(pv):.2f}
"""

    except:
        return None

    return None

# =========================================================
# EMBEDDING CACHE
# =========================================================

@lru_cache(maxsize=256)
def cached_query_embedding(query):

    return embedder.encode(
        [query],
        convert_to_numpy=True,
        normalize_embeddings=True
    ).astype("float32")

# =========================================================
# BUILD HISTORY
# =========================================================

def build_conversation_history(history):

    messages = []

    if not history:
        return messages

    try:

        for h in history[-20:]:

            # NEWER FORMAT

            if isinstance(h, dict):

                role = h.get("role")
                content = h.get("content")

                if role and content:

                    messages.append({
                        "role": role,
                        "content": str(content)[:2000]
                    })

            # OLDER FORMAT

            elif isinstance(h, (list, tuple)):

                if len(h) >= 2:

                    user_msg = h[0]
                    assistant_msg = h[1]

                    if user_msg:

                        messages.append({
                            "role": "user",
                            "content": str(user_msg)[:1500]
                        })

                    if assistant_msg:

                        messages.append({
                            "role": "assistant",
                            "content": str(assistant_msg)[:2000]
                        })

    except Exception as e:

        print("History parsing error:", e)

    return messages

# =========================================================
# QUERY REWRITE
# =========================================================

def rewrite_query_with_history(message, history):

    if not history:
        return message

    try:

        recent_context = []

        for h in history[-8:]:

            # NEW FORMAT

            if isinstance(h, dict):

                role = h.get("role")
                content = h.get("content")

                if role and content:

                    recent_context.append(
                        f"{role.capitalize()}: {content}"
                    )

            # OLD FORMAT

            elif isinstance(h, (list, tuple)) and len(h) >= 2:

                recent_context.append(
                    f"User: {h[0]}"
                )

                recent_context.append(
                    f"Assistant: {h[1]}"
                )

        convo = "\n".join(recent_context)

        rewrite_messages = [

            {
                "role": "system",
                "content": (
                    "Your ONLY task is rewriting follow-up queries.\n"
                    "Return ONLY the rewritten standalone query.\n"
                    "Do NOT answer the query.\n"
                    "Resolve references like:\n"
                    "- it\n"
                    "- this\n"
                    "- that\n"
                    "- they\n"
                    "- explain more\n"
                    "- simplify it\n"
                )
            },

            {
                "role": "user",
                "content": f"""
Conversation:
{convo}

Latest User Query:
{message}

Standalone rewritten query:
"""
            }
        ]

        rewritten = generate_response(
            rewrite_messages,
            max_tokens=40
        )

        rewritten = rewritten.strip()

        bad_patterns = [
            "i'm finbot",
            "assistant",
            "how can i help"
        ]

        for pattern in bad_patterns:

            if pattern in rewritten.lower():

                return message

        if len(rewritten) < 4:

            return message

        return rewritten

    except:

        return message

# =========================================================
# RETRIEVE CONTEXT
# =========================================================

def retrieve_context(query, top_k=5):

    global faiss_index

    if faiss_index is None:

        return ""

    q_emb = cached_query_embedding(query)

    scores, indices = faiss_index.search(
        q_emb,
        top_k
    )

    contexts = []

    for idx in indices[0]:

        if idx < len(doc_chunks):

            contexts.append(
                doc_chunks[idx][:450]
            )

    context = "\n\n".join(contexts)

    return context[:MAX_CONTEXT_CHARS]

# =========================================================
# MAIN CHATBOT
# =========================================================

def finance_chat(message, history):

    if history is None:
        history = []

    history = history[-MAX_HISTORY_TURNS:]

    # =====================================================
    # FINANCIAL CALCULATOR
    # =====================================================

    calc_result = solve_financial_query(message)

    if calc_result:

        messages = [

            {
                "role": "system",
                "content": SYSTEM_PROMPT
            },

            {
                "role": "user",
                "content": calc_result
            }
        ]

        response = generate_response(
            messages,
            max_tokens=300
        )

        return format_response(response)

    # =====================================================
    # QUERY REWRITE
    # =====================================================

    enhanced_query = rewrite_query_with_history(
        message,
        history
    )

    # =====================================================
    # RETRIEVAL
    # =====================================================

    context = retrieve_context(enhanced_query)

    # =====================================================
    # BUILD MESSAGES
    # =====================================================

    messages = [

        {
            "role": "system",
            "content": SYSTEM_PROMPT
        }
    ]

    messages.extend(
        build_conversation_history(history)
    )

    # =====================================================
    # USER PROMPT
    # =====================================================

    if context:

        user_prompt = f"""
Relevant Financial Knowledge:
{context}

User Question:
{message}

Resolved Intent:
{enhanced_query}

Instructions:
- Answer conversationally
- Maintain continuity
- Use retrieved context when useful
- Give detailed explanations
- Use examples if relevant
- Complete your answer fully
"""

    else:

        user_prompt = f"""
User Question:
{message}

Resolved Intent:
{enhanced_query}

Instructions:
- Answer conversationally
- Maintain conversation continuity
- Give detailed explanations
- Use examples if useful
"""

    messages.append({

        "role": "user",

        "content": user_prompt
    })

    response = generate_response(
        messages,
        max_tokens=500
    )

    return format_response(response)

# =========================================================
# DOCUMENT SYSTEM
# =========================================================

uploaded_chunks = []
uploaded_index = None

# =========================================================
# EXTRACT TEXT
# =========================================================

def extract_text(file):

    if file is None:
        return ""

    text = ""

    # PDF

    if file.name.endswith(".pdf"):

        with pdfplumber.open(file.name) as pdf:

            for page in pdf.pages:

                page_text = page.extract_text()

                if page_text:

                    text += page_text + "\n"

    # CSV

    elif file.name.endswith(".csv"):

        df = pd.read_csv(file.name)

        text += df.to_string()

    # TXT

    else:

        with open(file.name, "r", encoding="utf-8") as f:

            text += f.read()

    return re.sub(r"\s+", " ", text)

# =========================================================
# CHUNK TEXT
# =========================================================

def chunk_text(text, chunk_size=700, overlap=120):

    chunks = []

    start = 0

    while start < len(text):

        end = start + chunk_size

        chunks.append(text[start:end])

        start += chunk_size - overlap

    return chunks

# =========================================================
# PROCESS DOCUMENT
# =========================================================

def process_document(file):

    global uploaded_chunks
    global uploaded_index

    if file is None:

        return "❌ Please upload a document."

    text = extract_text(file)

    if not text.strip():

        return "❌ No readable text found."

    uploaded_chunks = chunk_text(text)

    embeddings = embedder.encode(
        uploaded_chunks,
        convert_to_numpy=True,
        normalize_embeddings=True,
        batch_size=32,
        show_progress_bar=False
    ).astype("float32")

    dimension = embeddings.shape[1]

    uploaded_index = faiss.IndexFlatIP(dimension)

    uploaded_index.add(embeddings)

    return f"""
✅ Document processed successfully

📄 Chunks created: {len(uploaded_chunks)}

🧠 You can now chat with your document.
"""

# =========================================================
# DOCUMENT CHAT
# =========================================================

def doc_chat(message, history):

    global uploaded_index

    if history is None:
        history = []

    if uploaded_index is None:

        return "Please upload and process a document first."

    q_emb = cached_query_embedding(message)

    scores, indices = uploaded_index.search(
        q_emb,
        5
    )

    contexts = []

    for idx in indices[0]:

        if idx < len(uploaded_chunks):

            contexts.append(
                uploaded_chunks[idx][:450]
            )

    context = "\n\n".join(contexts)

    messages = [

        {
            "role": "system",
            "content": SYSTEM_PROMPT
        }
    ]

    messages.extend(
        build_conversation_history(history)
    )

    messages.append({

        "role": "user",

        "content": f"""
Document Context:
{context}

Question:
{message}

Instructions:
- Use ONLY the document information
- Be conversational
- Maintain continuity
- Give detailed explanations
- Do not hallucinate
"""
    })

    response = generate_response(
        messages,
        max_tokens=500
    )

    return format_response(response)

# =========================================================
# MODERN UI CSS
# =========================================================

custom_css = """

body {
    background: linear-gradient(
        135deg,
        #0f172a,
        #111827
    );
}

.gradio-container {
    max-width: 1300px !important;
    margin: auto;
}

footer {
    visibility: hidden;
}

h1 {
    text-align: center;
    font-size: 48px !important;
    margin-bottom: 0px !important;

    background: linear-gradient(
        90deg,
        #38bdf8,
        #818cf8
    );

    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
}

.subtitle {
    text-align: center;
    color: #cbd5e1;
    font-size: 18px;
    margin-bottom: 20px;
}

.upload-card {

    background: rgba(30,41,59,0.6);

    border: 1px solid #334155;

    border-radius: 24px;

    padding: 25px;

    margin-bottom: 20px;
}

.upload-title {

    font-size: 26px;

    font-weight: bold;

    margin-bottom: 8px;
}

.upload-subtitle {

    color: #cbd5e1;

    margin-bottom: 15px;
}

.gr-chatbot {

    border-radius: 22px !important;

    border: 1px solid #334155 !important;
}

.gr-button {

    border-radius: 14px !important;

    font-size: 16px !important;
}

.gr-textbox textarea {

    font-size: 16px !important;
}

"""

# =========================================================
# UI
# =========================================================

with gr.Blocks(
    title="FinBot AI"
) as demo:

    gr.Markdown(
        """
# 💰 FinBot AI

<div class="subtitle">
Conversational Financial AI Assistant with RAG + Document Intelligence
</div>
"""
    )

    with gr.Tabs():

        # =================================================
        # MAIN CHATBOT
        # =================================================

        with gr.Tab("💬 Assistant"):

            gr.ChatInterface(

                fn=finance_chat,

                chatbot=gr.Chatbot(
                    height=650
                ),

                textbox=gr.Textbox(
                    placeholder="Ask finance questions naturally...",
                    container=False,
                    scale=7
                ),

                examples=[
                    "Explain inflation",
                    "What is EBITDA?",
                    "How does compound interest work?",
                    "Difference between NPV and IRR",
                    "Explain liquidity ratios"
                ]
            )

        # =================================================
        # DOCUMENT AI
        # =================================================

        with gr.Tab("📄 Document Intelligence"):

            gr.Markdown(
                """
<div class="upload-card">

<div class="upload-title">
📄 AI Financial Document Analyzer
</div>

<div class="upload-subtitle">
Upload PDFs, CSVs, or TXT files and chat intelligently with your documents.
</div>

</div>
"""
            )

            with gr.Row():

                file_input = gr.File(
                    label="Upload Document",
                    file_types=[
                        ".pdf",
                        ".csv",
                        ".txt"
                    ]
                )

                upload_btn = gr.Button(
                    "⚡ Process Document",
                    variant="primary"
                )

            status = gr.Textbox(
                label="Processing Status",
                lines=5
            )

            upload_btn.click(
                process_document,
                inputs=file_input,
                outputs=status
            )

            gr.ChatInterface(

                fn=doc_chat,

                chatbot=gr.Chatbot(
                    height=650
                ),

                textbox=gr.Textbox(
                    placeholder="Ask questions about your uploaded document...",
                    container=False
                ),

                examples=[
                    "Summarize this report",
                    "Analyze profitability",
                    "What risks are mentioned?",
                    "Explain liabilities",
                    "Give key insights"
                ]
            )

# =========================================================
# QUEUE
# =========================================================

demo.queue(
    max_size=20,
    default_concurrency_limit=2
)

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

if __name__ == "__main__":

    demo.launch(
        css=custom_css,
        theme=gr.themes.Soft(
            primary_hue="blue",
            secondary_hue="slate"
        )
    )