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
from langgraph.graph import StateGraph
from typing import TypedDict, List

# Set up Streamlit page
st.set_page_config(page_title="Financial Insights Chatbot", page_icon="📊", layout="wide")
GROQ_API_KEY = "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"
ALPHA_VANTAGE_API_KEY = "BUJB69F41Y4ESEGX"
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")

temp_llm = ChatGroq(temperature=0, model="llama3-70b-8192", api_key=GROQ_API_KEY)
st.success("✅ LLM initialized successfully. Using llama3-70b-8192")

embedding_model = SentenceTransformer("baconnier/Finance2_embedding_small_en-V1.5", device=device)
text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)

class AgentState(TypedDict):
    extracted_data: str
    compliance_issues: str
    risk_analysis: str
    report: str
    audit_log: List[str]
    approval_status: str

def extract_data(state: AgentState) -> AgentState:
    """Extracts financial data from PDFs or APIs."""
    if state.get("extracted_data"):
        return state
    return {**state, "extracted_data": "Extracted financial data here."}

def validate_compliance(state: AgentState) -> AgentState:
    """Validates financial data against compliance rules."""
    return {**state, "compliance_issues": "No issues found."}

def assess_risk(state: AgentState) -> AgentState:
    """Analyzes potential financial risks."""
    return {**state, "risk_analysis": "Low risk identified."}

def generate_report(state: AgentState) -> AgentState:
    """Generates a financial compliance report."""
    report = f"Report:\nData: {state['extracted_data']}\nCompliance: {state['compliance_issues']}\nRisk: {state['risk_analysis']}"
    return {**state, "report": report}

def log_audit(state: AgentState) -> AgentState:
    """Maintains an audit trail of changes."""
    return {**state, "audit_log": state.get("audit_log", []) + ["Report generated"]}

def approve_report(state: AgentState) -> AgentState:
    """Sends the report for approval."""
    return {**state, "approval_status": "Approved"}

def display_dashboard(state: AgentState):
    """Shows the final insights on Streamlit UI."""
    st.markdown("### 📊 Final Dashboard")
    st.write(state)

workflow = StateGraph(AgentState)
workflow.add_node("extract_data", extract_data)
workflow.add_node("validate_compliance", validate_compliance)
workflow.add_node("assess_risk", assess_risk)
workflow.add_node("generate_report", generate_report)
workflow.add_node("log_audit", log_audit)
workflow.add_node("approve_report", approve_report)
workflow.add_edge("extract_data", "validate_compliance")
workflow.add_edge("validate_compliance", "assess_risk")
workflow.add_edge("assess_risk", "generate_report")
workflow.add_edge("generate_report", "log_audit")
workflow.add_edge("log_audit", "approve_report")
workflow.set_entry_point("extract_data")
workflow.set_termination_nodes(["approve_report"])

if st.button("Run Multi-Agent Workflow"):
    output = workflow.run({})
    display_dashboard(output)