import streamlit as st
import pymupdf  
import re
import traceback
import faiss
import numpy as np
import requests
from rank_bm25 import BM25Okapi
from sentence_transformers import SentenceTransformer
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_groq import ChatGroq
import torch
import os 

st.set_page_config(page_title="Financial Insights Chatbot", page_icon="📊", layout="wide")

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

GROQ_API_KEY = os.getenv("GROQ_API_KEY")
ALPHA_VANTAGE_API_KEY = os.getenv("ALPHA_VANTAGE_API_KEY")

try:
    llm = ChatGroq(temperature=0, model="llama3-70b-8192", api_key=GROQ_API_KEY)
    st.success("✅ LLM initialized successfully. Using llama3-70b-8192")
except Exception as e:
    st.error("❌ Failed to initialize Groq LLM.")
    traceback.print_exc()

embedding_model = SentenceTransformer("baconnier/Finance2_embedding_small_en-V1.5", device=device)

text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)

def fetch_financial_data(company_ticker):
    if not company_ticker:
        return "No ticker symbol provided. Please enter a valid company ticker."

    try:
        overview_url = f"https://www.alphavantage.co/query?function=OVERVIEW&symbol={company_ticker}&apikey={ALPHA_VANTAGE_API_KEY}"
        overview_response = requests.get(overview_url)

        if overview_response.status_code == 200:
            overview_data = overview_response.json()
            market_cap = overview_data.get("MarketCapitalization", "N/A")
        else:
            return "Error fetching company overview."

        income_url = f"https://www.alphavantage.co/query?function=INCOME_STATEMENT&symbol={company_ticker}&apikey={ALPHA_VANTAGE_API_KEY}"
        income_response = requests.get(income_url)

        if income_response.status_code == 200:
            income_data = income_response.json()
            annual_reports = income_data.get("annualReports", [])
            revenue = annual_reports[0].get("totalRevenue", "N/A") if annual_reports else "N/A"
        else:
            return "Error fetching income statement."

        return f"Market Cap: ${market_cap}\nTotal Revenue: ${revenue}"

    except Exception as e:
        traceback.print_exc()
        return "Error fetching financial data."

def extract_and_embed_text(pdf_file):
    """Processes PDFs and generates embeddings with GPU acceleration using pymupdf."""
    try:
        docs, tokenized_texts = [], []

        with pymupdf.open(stream=pdf_file.read(), filetype="pdf") as doc:
            full_text = "\n".join(page.get_text("text") for page in doc)
            chunks = text_splitter.split_text(full_text)
            for chunk in chunks:
                docs.append(chunk)
                tokenized_texts.append(chunk.split())

        embeddings = embedding_model.encode(docs, batch_size=64, convert_to_numpy=True, normalize_embeddings=True)

        embedding_dim = embeddings.shape[1]
        index = faiss.IndexHNSWFlat(embedding_dim, 32)
        index.add(embeddings)

        bm25 = BM25Okapi(tokenized_texts)

        return docs, embeddings, index, bm25
    except Exception as e:
        traceback.print_exc()
        return [], [], None, None

def retrieve_relevant_docs(user_query, docs, index, bm25):
    """Hybrid search using FAISS cosine similarity & BM25 keyword retrieval."""
    query_embedding = embedding_model.encode(user_query, convert_to_numpy=True, normalize_embeddings=True)
    _, faiss_indices = index.search(np.array([query_embedding]), 8)
    bm25_scores = bm25.get_scores(user_query.split())
    bm25_indices = np.argsort(bm25_scores)[::-1][:8]
    combined_indices = list(set(faiss_indices[0]) | set(bm25_indices))

    return [docs[i] for i in combined_indices[:3]]

def generate_response(user_query, pdf_ticker, ai_ticker, mode, uploaded_file):
    try:
        if mode == "📄 PDF Upload Mode":
            docs, embeddings, index, bm25 = extract_and_embed_text(uploaded_file)
            if not docs:
                return "❌ Error extracting text from PDF."

            retrieved_docs = retrieve_relevant_docs(user_query, docs, index, bm25)
            context = "\n\n".join(retrieved_docs)
            prompt = f"Summarize the key financial insights for {pdf_ticker} from this document:\n\n{context}"

        elif mode == "🌍 Live Data Mode":
            financial_info = fetch_financial_data(ai_ticker)
            prompt = f"Analyze the financial status of {ai_ticker} based on:\n{financial_info}\n\nUser Query: {user_query}"
        else:
            return "Invalid mode selected."

        response = llm.invoke(prompt)
        return response.content
    except Exception as e:
        traceback.print_exc()
        return "Error generating response."

st.markdown(
    "<h1 style='text-align: center; color: #4CAF50;'>📄 FinQuery RAG Chatbot</h1>", 
    unsafe_allow_html=True
)
st.markdown(
    "<h5 style='text-align: center; color: #666;'>Analyze financial reports or fetch live financial data effortlessly!</h5>", 
    unsafe_allow_html=True
)

col1, col2 = st.columns(2)

with col1:
    st.markdown("### 🏢 **Choose Your Analysis Mode**")
    mode = st.radio("", ["📄 PDF Upload Mode", "🌍 Live Data Mode"], horizontal=True)

with col2:
    st.markdown("### 🔎 **Enter Your Query**")
    user_query = st.text_input("💬 What financial insights are you looking for?")

st.markdown("---")
if mode == "📄 PDF Upload Mode":
    st.markdown("### 📂 Upload Your Financial Report")
    uploaded_file = st.file_uploader("🔼 Upload PDF (Only for PDF Mode)", type=["pdf"])
    pdf_ticker = st.text_input("🏢 Enter Company Ticker for PDF Insights", placeholder="e.g., INFY, TCS")
    ai_ticker = None
else:
    st.markdown("### 🌍 Live Market Data")
    ai_ticker = st.text_input("🏢 Enter Company Ticker for AI Insights", placeholder="e.g., AAPL, MSFT")
    uploaded_file = None
    pdf_ticker = None

if st.button("Analyze Now"):
    if mode == "📄 PDF Upload Mode" and (not uploaded_file or not pdf_ticker):
        st.error("❌ Please upload a PDF and enter a company ticker for insights.")
    elif mode == "🌍 Live Data Mode" and not ai_ticker:
        st.error("❌ Please enter a valid company ticker for AI insights.")
    else:
        with st.spinner("🔍 Your Query is Processing, this can take up to 5 - 7 minutes ⏳"):
            response = generate_response(user_query, pdf_ticker, ai_ticker, mode, uploaded_file)
            st.markdown("---")
            st.markdown("<h3 style='color: #4CAF50;'>💡 AI Response</h3>", unsafe_allow_html=True)
            st.write(response)

st.markdown("---")