import gradio as gr
import groq
import os
import tempfile
import uuid
import yfinance as yf
import pandas as pd
import plotly.graph_objects as go
from dotenv import load_dotenv
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.vectorstores import FAISS
from langchain.embeddings import HuggingFaceEmbeddings
import fitz  # PyMuPDF
import base64
from PIL import Image
import io
import requests
import json

# Load environment variables
load_dotenv()
client = groq.Client(api_key=os.getenv("GROQ_LEGAL_API_KEY"))
embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")

# Directory to store FAISS indexes
FAISS_INDEX_DIR = "faiss_indexes_finance"
if not os.path.exists(FAISS_INDEX_DIR):
    os.makedirs(FAISS_INDEX_DIR)

# Dictionary to store user-specific vectorstores
user_vectorstores = {}

# Custom CSS for Finance theme
custom_css = """
:root {
    --primary-color: #FFD700;  /* Gold */
    --secondary-color: #008000;  /* Dark Green */
    --light-background: #F0FFF0;  /* Honeydew */
    --dark-text: #333333;
    --white: #FFFFFF;
    --border-color: #E5E7EB;
}
body { background-color: var(--light-background); font-family: 'Inter', sans-serif; }
.container { max-width: 1200px !important; margin: 0 auto !important; padding: 10px; }
.header { background-color: var(--white); border-bottom: 2px solid var(--border-color); padding: 15px 0; margin-bottom: 20px; border-radius: 12px 12px 0 0; box-shadow: 0 2px 4px rgba(0,0,0,0.05); }
.header-title { color: var(--secondary-color); font-size: 1.8rem; font-weight: 700; text-align: center; }
.header-subtitle { color: var(--dark-text); font-size: 1rem; text-align: center; margin-top: 5px; }
.chat-container { border-radius: 12px !important; box-shadow: 0 4px 6px rgba(0,0,0,0.1) !important; background-color: var(--white) !important; border: 1px solid var(--border-color) !important; min-height: 500px; }
.message-user { background-color: var(--primary-color) !important; color: var(--dark-text) !important; border-radius: 18px 18px 4px 18px !important; padding: 12px 16px !important; margin-left: auto !important; max-width: 80% !important; }
.message-bot { background-color: #F0F0F0 !important; color: var(--dark-text) !important; border-radius: 18px 18px 18px 4px !important; padding: 12px 16px !important; margin-right: auto !important; max-width: 80% !important; }
.input-area { background-color: var(--white) !important; border-top: 1px solid var(--border-color) !important; padding: 12px !important; border-radius: 0 0 12px 12px !important; }
.input-box { border: 1px solid var(--border-color) !important; border-radius: 24px !important; padding: 12px 16px !important; box-shadow: 0 2px 4px rgba(0,0,0,0.05) !important; }
.send-btn { background-color: var(--secondary-color) !important; border-radius: 24px !important; color: var(--white) !important; padding: 10px 20px !important; font-weight: 500 !important; }
.clear-btn { background-color: #F0F0F0 !important; border: 1px solid var(--border-color) !important; border-radius: 24px !important; color: var(--dark-text) !important; padding: 8px 16px !important; font-weight: 500 !important; }
.pdf-viewer-container { border-radius: 12px !important; box-shadow: 0 4px 6px rgba(0,0,0,0.1) !important; background-color: var(--white) !important; border: 1px solid var(--border-color) !important; padding: 20px; }
.pdf-viewer-image { max-width: 100%; height: auto; border: 1px solid var(--border-color); border-radius: 12px; box-shadow: 0 2px 4px rgba(0,0,0,0.05); }
.stats-box { background-color: #E6F2E6; padding: 10px; border-radius: 8px; margin-top: 10px; }
.tool-container { background-color: var(--white); border-radius: 12px; box-shadow: 0 4px 6px rgba(0,0,0,0.1); padding: 15px; margin-bottom: 20px; }
.tool-title { color: var(--secondary-color); font-size: 1.2rem; font-weight: 600; margin-bottom: 10px; }
.chart-container { height: 400px; width: 100%; border-radius: 8px; overflow: hidden; }
"""

# Function to process PDF files (unchanged)
def process_pdf(pdf_file):
    if pdf_file is None:
        return None, "No file uploaded", {"page_images": [], "total_pages": 0, "total_words": 0}
    try:
        session_id = str(uuid.uuid4())
        with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as temp_file:
            temp_file.write(pdf_file)
            pdf_path = temp_file.name
        
        doc = fitz.open(pdf_path)
        texts = [page.get_text() for page in doc]
        page_images = []
        for page in doc:
            pix = page.get_pixmap()
            img_bytes = pix.tobytes("png")
            img_base64 = base64.b64encode(img_bytes).decode("utf-8")
            page_images.append(img_base64)
        total_pages = len(doc)
        total_words = sum(len(text.split()) for text in texts)
        doc.close()

        text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
        chunks = text_splitter.create_documents(texts)
        vectorstore = FAISS.from_documents(chunks, embeddings)
        index_path = os.path.join(FAISS_INDEX_DIR, session_id)
        vectorstore.save_local(index_path)
        user_vectorstores[session_id] = vectorstore

        os.unlink(pdf_path)
        pdf_state = {"page_images": page_images, "total_pages": total_pages, "total_words": total_words}
        return session_id, f"✅ Successfully processed {len(chunks)} text chunks from your PDF", pdf_state
    except Exception as e:
        if "pdf_path" in locals() and os.path.exists(pdf_path):
            os.unlink(pdf_path)
        return None, f"Error processing PDF: {str(e)}", {"page_images": [], "total_pages": 0, "total_words": 0}

# Function to generate chatbot responses with Finance theme
def generate_response(message, session_id, model_name, history):
    if not message:
        return history
    try:
        context = ""
        if session_id and session_id in user_vectorstores:
            vectorstore = user_vectorstores[session_id]
            docs = vectorstore.similarity_search(message, k=3)
            if docs:
                context = "\n\nRelevant information from uploaded PDF:\n" + "\n".join(f"- {doc.page_content}" for doc in docs)
        
        # Check if it's a stock ticker query
        if message.startswith("$") and len(message) > 1 and len(message) <= 6:
            ticker = message[1:].upper()
            try:
                stock_data = get_stock_data(ticker)
                response = f"**Stock Information for {ticker}**\n\n"
                response += f"Current Price: ${stock_data['current_price']}\n"
                response += f"52-Week High: ${stock_data['52wk_high']}\n"
                response += f"Market Cap: ${stock_data['market_cap']:,}\n"
                response += f"P/E Ratio: {stock_data['pe_ratio']}\n"
                response += f"More data available in the Stock Analysis tab."
                history.append((message, response))
                return history
            except Exception as e:
                history.append((message, f"Error retrieving stock data for {ticker}: {str(e)}"))
                return history
                
        system_prompt = "You are a financial assistant specializing in analyzing financial reports, statements, and market trends."
        system_prompt += " You can help with stock market information, financial terminology, ratio analysis, and investment concepts."
        if context:
            system_prompt += " Use the following context to answer the question if relevant: " + context
        
        completion = client.chat.completions.create(
            model=model_name,
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": message}
            ],
            temperature=0.7,
            max_tokens=1024
        )
        response = completion.choices[0].message.content
        history.append((message, response))
        return history
    except Exception as e:
        history.append((message, f"Error generating response: {str(e)}"))
        return history

# Functions to update PDF viewer (unchanged)
def update_pdf_viewer(pdf_state):
    if not pdf_state["total_pages"]:
        return 0, None, "No PDF uploaded yet"
    try:
        img_data = base64.b64decode(pdf_state["page_images"][0])
        img = Image.open(io.BytesIO(img_data))
        return pdf_state["total_pages"], img, f"**Total Pages:** {pdf_state['total_pages']}\n**Total Words:** {pdf_state['total_words']}"
    except Exception as e:
        print(f"Error decoding image: {e}")
        return 0, None, "Error displaying PDF"

def update_image(page_num, pdf_state):
    if not pdf_state["total_pages"] or page_num < 1 or page_num > pdf_state["total_pages"]:
        return None
    try:
        img_data = base64.b64decode(pdf_state["page_images"][page_num - 1])
        img = Image.open(io.BytesIO(img_data))
        return img
    except Exception as e:
        print(f"Error decoding image: {e}")
        return None

# New Finance-specific tools
def get_stock_data(ticker):
    """Tool to fetch latest stock data for a given ticker"""
    try:
        stock = yf.Ticker(ticker)
        info = stock.info
        return {
            "current_price": info.get("currentPrice", info.get("regularMarketPrice", "N/A")),
            "52wk_high": info.get("fiftyTwoWeekHigh", "N/A"),
            "market_cap": info.get("marketCap", "N/A"),
            "pe_ratio": info.get("trailingPE", "N/A"),
            "dividend_yield": info.get("dividendYield", "N/A"),
            "beta": info.get("beta", "N/A"),
            "average_volume": info.get("averageVolume", "N/A")
        }
    except Exception as e:
        print(f"Error fetching stock data: {e}")
        raise e

def get_stock_history(ticker, period="1y"):
    """Get historical data for charting"""
    try:
        stock = yf.Ticker(ticker)
        hist = stock.history(period=period)
        return hist
    except Exception as e:
        print(f"Error fetching stock history: {e}")
        return pd.DataFrame()

def get_fred_data(indicator):
    """Get economic data from FRED API"""
    api_key = os.getenv("FRED_API_KEY", "")
    if not api_key:
        return "FRED API key not configured"
    
    base_url = "https://api.stlouisfed.org/fred/series/observations"
    params = {
        "series_id": indicator,
        "api_key": api_key,
        "file_type": "json",
        "sort_order": "desc",
        "limit": 100
    }
    
    try:
        response = requests.get(base_url, params=params)
        data = response.json()
        return data.get("observations", [])
    except Exception as e:
        print(f"Error fetching FRED data: {e}")
        return []

def create_stock_chart(ticker, period="1y"):
    """Create an interactive stock chart using Plotly"""
    try:
        df = get_stock_history(ticker, period)
        if df.empty:
            return None
            
        fig = go.Figure()
        
        # Add candlestick chart
        fig.add_trace(
            go.Candlestick(
                x=df.index,
                open=df['Open'],
                high=df['High'],
                low=df['Low'],
                close=df['Close'],
                name=ticker
            )
        )
        
        # Add volume as bar chart on secondary y-axis
        fig.add_trace(
            go.Bar(
                x=df.index,
                y=df['Volume'],
                name='Volume',
                marker_color='rgba(0, 128, 0, 0.3)',
                yaxis='y2'
            )
        )
        
        # Update layout for dual y-axis
        fig.update_layout(
            title=f'{ticker} Stock Price',
            yaxis_title='Price (USD)',
            xaxis_title='Date',
            template='plotly_white',
            yaxis=dict(
                domain=[0.3, 1.0]
            ),
            yaxis2=dict(
                domain=[0, 0.2],
                title='Volume'
            ),
            legend=dict(
                orientation="h",
                yanchor="bottom",
                y=1.02,
                xanchor="right",
                x=1
            ),
            height=500
        )
        
        return fig
    except Exception as e:
        print(f"Error creating stock chart: {e}")
        return None

def analyze_ticker(ticker_input, period):
    """Process the ticker input and return analysis"""
    if not ticker_input:
        return None, "Please enter a valid ticker symbol", None
    
    ticker = ticker_input.strip().upper()
    if ticker.startswith("$"):
        ticker = ticker[1:]
    
    try:
        stock_data = get_stock_data(ticker)
        chart = create_stock_chart(ticker, period)
        
        # Create a formatted summary
        summary = f"""
### {ticker} Analysis
**Current Price:** ${stock_data['current_price']}  
**52-Week High:** ${stock_data['52wk_high']}  
**Market Cap:** ${stock_data['market_cap']:,}  
**P/E Ratio:** {stock_data['pe_ratio']}  
**Dividend Yield:** {stock_data['dividend_yield'] * 100 if stock_data['dividend_yield'] != 'N/A' else 'N/A'}%  
**Beta:** {stock_data['beta']}  
**Avg Volume:** {stock_data['average_volume']:,}  
        """
        
        return chart, summary, ticker
    except Exception as e:
        return None, f"Error analyzing ticker {ticker}: {str(e)}", None

# Gradio interface
with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
    current_session_id = gr.State(None)
    pdf_state = gr.State({"page_images": [], "total_pages": 0, "total_words": 0})
    current_ticker = gr.State(None)
    
    gr.HTML("""
    <div class="header">
        <div class="header-title">Fin-Vision</div>
        <div class="header-subtitle">Analyze financial documents with Groq's LLM API.</div>
    </div>
    """)
    
    with gr.Row(elem_classes="container"):
        with gr.Column(scale=1, min_width=300):
            pdf_file = gr.File(label="Upload PDF Document", file_types=[".pdf"], type="binary")
            upload_button = gr.Button("Process PDF", variant="primary")
            pdf_status = gr.Markdown("No PDF uploaded yet")
            model_dropdown = gr.Dropdown(
                choices=["llama3-70b-8192", "llama3-8b-8192", "mixtral-8x7b-32768", "gemma-7b-it"],
                value="llama3-70b-8192",
                label="Select Groq Model"
            )
            
            # Finance Tools Section
            gr.Markdown("### Financial Tools", elem_classes="tool-title")
            with gr.Group(elem_classes="tool-container"):
                with gr.Tabs():
                    with gr.TabItem("Stock Analysis"):
                        ticker_input = gr.Textbox(label="Enter Ticker Symbol (e.g., AAPL)", placeholder="AAPL")
                        period_dropdown = gr.Dropdown(
                            choices=["1mo", "3mo", "6mo", "1y", "2y", "5y", "max"],
                            value="1y",
                            label="Time Period"
                        )
                        analyze_button = gr.Button("Analyze Stock")
            
        with gr.Column(scale=2, min_width=600):
            with gr.Tabs():
                with gr.TabItem("PDF Viewer"):
                    with gr.Column(elem_classes="pdf-viewer-container"):
                        page_slider = gr.Slider(minimum=1, maximum=1, step=1, label="Page Number", value=1)
                        pdf_image = gr.Image(label="PDF Page", type="pil", elem_classes="pdf-viewer-image")
                        stats_display = gr.Markdown("No PDF uploaded yet", elem_classes="stats-box")
                
                with gr.TabItem("Stock Analysis"):
                    with gr.Column(elem_classes="pdf-viewer-container"):
                        stock_chart = gr.Plot(label="Stock Price Chart", elem_classes="chart-container")
                        stock_summary = gr.Markdown("Enter a ticker symbol to see analysis")
    
    with gr.Row(elem_classes="container"):
        with gr.Column(scale=2, min_width=600):
            chatbot = gr.Chatbot(height=500, bubble_full_width=False, show_copy_button=True, elem_classes="chat-container")
            with gr.Row():
                msg = gr.Textbox(show_label=False, placeholder="Ask about your financial document or type $TICKER for stock info...", scale=5)
                send_btn = gr.Button("Send", scale=1)
            clear_btn = gr.Button("Clear Conversation")
    
    # Event Handlers
    upload_button.click(
        process_pdf,
        inputs=[pdf_file],
        outputs=[current_session_id, pdf_status, pdf_state]
    ).then(
        update_pdf_viewer,
        inputs=[pdf_state],
        outputs=[page_slider, pdf_image, stats_display]
    )
    
    msg.submit(
        generate_response,
        inputs=[msg, current_session_id, model_dropdown, chatbot],
        outputs=[chatbot]
    ).then(lambda: "", None, [msg])
    
    send_btn.click(
        generate_response,
        inputs=[msg, current_session_id, model_dropdown, chatbot],
        outputs=[chatbot]
    ).then(lambda: "", None, [msg])
    
    clear_btn.click(
        lambda: ([], None, "No PDF uploaded yet", {"page_images": [], "total_pages": 0, "total_words": 0}, 0, None, "No PDF uploaded yet", None),
        None,
        [chatbot, current_session_id, pdf_status, pdf_state, page_slider, pdf_image, stats_display, current_ticker]
    )
    
    page_slider.change(
        update_image,
        inputs=[page_slider, pdf_state],
        outputs=[pdf_image]
    )
    
    # Stock analysis handler
    analyze_button.click(
        analyze_ticker,
        inputs=[ticker_input, period_dropdown],
        outputs=[stock_chart, stock_summary, current_ticker]
    )

# Add footer with attribution
gr.HTML("""
<div style="text-align: center; margin-top: 20px; padding: 10px; color: #666; font-size: 0.8rem; border-top: 1px solid #eee;">
    Created by Calvin Allen Crawford
</div>
""")

# Launch the app
if __name__ == "__main__":
    demo.launch()