import streamlit as st
import requests
import pymupdf 
import traceback
from sentence_transformers import SentenceTransformer
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_groq import ChatGroq

ALPHA_VANTAGE_API_KEY = st.secrets["ALPHA_VANTAGE_API_KEY"]
GROQ_API_KEY = st.secrets["GROQ_API_KEY"]
embedding_model = SentenceTransformer("all-MiniLM-L6-v2")

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

def extract_text_from_pdf(uploaded_file, max_length=5000):
    try:
        doc = pymupdf.open(stream=uploaded_file.read(), filetype="pdf")  
        full_text = "".join(page.get_text() for page in doc)
        text_splitter = RecursiveCharacterTextSplitter(chunk_size=max_length, chunk_overlap=200)
        chunks = text_splitter.split_text(full_text)

        return chunks 
    except Exception as e:
        st.error("❌ Failed to extract text from PDF.")
        traceback.print_exc()
        return ["Error extracting text."]

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:
            st.error(f"❌ Failed to fetch company overview. Status Code: {overview_response.status_code}")
            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:
            st.error(f"❌ Failed to fetch income statement. Status Code: {income_response.status_code}")
            return "Error fetching income statement."

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

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

def generate_response(user_query, company_ticker, mode, uploaded_file):
    try:
        if mode == "PDF Upload Mode":
            chunks = extract_text_from_pdf(uploaded_file)
            chunked_summary = "\n\n".join(chunks[:3]) 
            prompt = f"Summarize the key financial insights from this document:\n\n{chunked_summary}"
        elif mode == "Live Data Mode":
            financial_info = fetch_financial_data(company_ticker)
            prompt = f"Analyze the financial status of {company_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:
        st.error("❌ Failed to generate AI response.")
        traceback.print_exc()
        return "Error generating response."

st.title("📊 AI-Powered Financial Insights Chatbot")
st.write("Upload financial reports or fetch live financial data to get AI-driven insights.")


user_query = st.text_input("Enter your query:")
company_ticker = st.text_input("Enter company ticker symbol (optional):")
mode = st.radio("Select Mode:", ["PDF Upload Mode", "Live Data Mode"])
uploaded_file = st.file_uploader("Upload PDF (Only for PDF Mode)", type=["pdf"])

if st.button("Get Insights"):
    if mode == "PDF Upload Mode" and not uploaded_file:
        st.error("❌ Please upload a PDF file.")
    else:
        with st.spinner("Processing... ⏳"):
            response = generate_response(user_query, company_ticker, mode, uploaded_file)
            st.subheader("💡 AI Response")
            st.write(response)