import nest_asyncio
nest_asyncio.apply()

import streamlit as st
from dotenv import load_dotenv
import io

from langchain_community.vectorstores import FAISS
from langchain_google_genai import ChatGoogleGenerativeAI, GoogleGenerativeAIEmbeddings
from langchain.prompts import ChatPromptTemplate
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain.tools import Tool
from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_core.messages import HumanMessage, ToolMessage
from pydantic import BaseModel, Field
from google.generativeai.types import HarmCategory, HarmBlockThreshold
from langchain.docstore.document import Document
import pdfplumber
import pandas as pd
import re

from langchain.chains import create_history_aware_retriever, create_retrieval_chain
from langchain_core.prompts import MessagesPlaceholder

# loading api keys
load_dotenv()

# main app interface
st.set_page_config(page_title="AI IPO Analyst", page_icon="📈")
st.title("📈 AI-Powered IPO Analysis Agent")
st.write("Upload a company's DRHP (Draft Red Herring Prospectus) in PDF format to begin the analysis.")

if 'analysis_report' not in st.session_state:
    st.session_state.analysis_report =""

def clean_value(value):
    """Helper function to clean individual cell values."""
    if value is None:
        return ""
    # Remove newlines and extra whitespace
    return str(value).replace('\n', ' ').strip()

# logic
@st.cache_data #caching to avoid rerunning on every interaction
def find_financial_section_pages(pdf_bytes):
    """
    Step 1: The Librarian.
    Reads the Table of Contents to find the page range for the financial statements.
    """
    st.write("  -> Reading Table of Contents to locate financial section...")
    toc_pages = []
    
    toc_pattern = re.compile(r'^(.*?)\s*([.\s]*)\s*(\d+)$')
    financial_section_keyword = re.compile(r'(?i)financial\s+(statements|information)')
    
    start_page = -1

    with pdfplumber.open(io.BytesIO(pdf_bytes)) as pdf:
        for i in range(min(20, len(pdf.pages))):
            page_text = pdf.pages[i].extract_text()
            if not page_text: continue
            
            for line in page_text.split('\n'):
                if financial_section_keyword.search(line):
                    match = toc_pattern.search(line)
                    if match:
                        try:
                            
                            page_num = int(match.group(3)) - 1
                            start_page = page_num
                            st.write(f"  ✅ Found 'Financial Statements' section in ToC, starting around page {page_num + 1}.")
                            
                            return range(start_page, min(start_page + 40, len(pdf.pages)))
                        except ValueError:
                            continue
    
    st.warning("Could not find 'Financial Statements' in the Table of Contents. Falling back to full document search.")
    return None

def find_financial_statements(pdf_bytes, page_range=None):
    """
    Step 2: The Detective.
    Searches within a given page range and uses a scoring system to find the best tables.
    If no page range is provided, it searches the whole document as a fallback.
    """
    
    anchor_keywords = {
        "balance_sheet": ["total equity and liabilities", "total assets"],
        "profit_and_loss": ["revenue from operations", "profit for the period"],
        "cash_flow": ["net cash generated from operating activities", "cash flow from operating activities"],
    }
    statement_keywords = {
        "balance_sheet": re.compile(r'(?i)BALANCE\s+SHEET'),
        "profit_and_loss": re.compile(r'(?i)PROFIT\s+&\s+LOSS|PROFIT\s+AND\s+LOSS'),
        "cash_flow": re.compile(r'(?i)CASH\s+FLOW\s+STATEMENT')
    }
    extracted_data = {}
    
    with pdfplumber.open(io.BytesIO(pdf_bytes)) as pdf:
        search_range = page_range if page_range else range(len(pdf.pages))

        for statement_name, keyword_regex in statement_keywords.items():
            best_table_overall = None
            highest_score_overall = -1

            for page_num in search_range:
                page = pdf.pages[page_num]
                page_text = page.extract_text(x_tolerance=2)
                if not page_text or not keyword_regex.search(page_text):
                    continue

                st.write(f"  -> Evaluating tables for '{statement_name}' on page {page_num + 1}...")
                tables = page.extract_tables()
                if not tables: continue

                for table_data in tables:
                    if not table_data or len(table_data) < 3: continue
                    current_score = 0
                    header = [str(h).replace('\n', ' ').strip() for h in table_data[0]]
                    year_pattern = re.compile(r'\b(20\d{2}|FY\d{2})\b')
                    years_found = len(year_pattern.findall(" ".join(header)))
                    current_score += years_found * 5
                    current_score += len(table_data)
                    first_column_text = " ".join([str(row[0]).lower() for row in table_data[1:] if row and len(row)>0])
                    for keyword in anchor_keywords[statement_name]:
                        if keyword in first_column_text:
                            current_score += 10
                    if current_score > highest_score_overall:
                        highest_score_overall = current_score
                        best_table_overall = table_data
            
            if best_table_overall:
                header = [str(h).replace('\n', ' ').strip() for h in best_table_overall[0]]
                df = pd.DataFrame(best_table_overall[1:], columns=header)
                df = df.dropna(how='all')
                extracted_data[statement_name] = df
                st.write(f"  ✅ Selected best table for '{statement_name}' (Score: {highest_score_overall}).")

    if not extracted_data:
        st.warning("Could not find and validate any standard financial statements.")
        return None
        
    return extracted_data

@st.cache_resource
def get_agent_components(_pdf_bytes):
    st.info("Step 1: Extracting financial tables from PDF...")
    
    
    page_range = find_financial_section_pages(_pdf_bytes)
    financial_data_dfs = find_financial_statements(_pdf_bytes, page_range)

    if not financial_data_dfs: return None, None, None

@st.cache_resource
def get_agent(_pdf_bytes, _company_name):
    """
    This is the new, fully functional get_agent function.
    It performs the entire pipeline: extract, ingest and agent setup.
    The underscores in the arguments (_pdf_bytes, _company_name) tell Streamlit's cachethis this function's output depends on the content of these arguments.
    """

    # extracting
    st.info("Step 1: Extracting financial tables from PDF...")
    financial_data_dfs = find_financial_statements(_pdf_bytes)
    if not financial_data_dfs:
        return None, None
    
    # ingesting
    st.info("Step 2: Preparing agent's memory (Vector Store)...")
    documents = []
    for statement_name, df in financial_data_dfs.items():
        content = f"This is the {statement_name.replace('_', ' ').title()}.\n\n{df.to_markdown(index=False)}"
        doc = Document(page_content=content, metadata={"source": statement_name})
        documents.append(doc)

    embeddings = GoogleGenerativeAIEmbeddings(model="models/embedding-001")
    vector_store = FAISS.from_documents(documents, embedding=embeddings)
    retriever = vector_store.as_retriever()
    st.write("✅ Agent's memory is ready.")

    # agent and tool setup
    st.info("Step 3: Initializing AI Agent and Tools...")
    llm = ChatGoogleGenerativeAI(
        model="gemini-2.5-pro",
        temperature=0,
        safety_settings={
            HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_NONE
        },
    )
    
    document_qa_prompt = ChatPromptTemplate.from_template("Answer based on context:\n{context}\nQuestion: {input}")
    document_chain = create_stuff_documents_chain(llm, document_qa_prompt)

    class FinancialSearchInput(BaseModel):
        query: str = Field(description="A detailed question to ask the financial statements.")
    
    financial_statement_tool = Tool(
        name="financial_statement_search",
        func=lambda query: document_chain.invoke({"input": query, "context": retriever.invoke(query)}),
        description="Search for info within the company's financial statements.",
        args_schema=FinancialSearchInput
    )

    web_search_tool = TavilySearchResults(name="web_search")

    tools = [financial_statement_tool, web_search_tool]

    llm_with_tools = llm.bind_tools(tools)
    st.write("✅ Agent and tools are ready.")

    return llm_with_tools, tools, retriever

def get_conversational_chain(retriever):
    """Creates the conversational RAG chain for the chat feature."""
    llm = ChatGoogleGenerativeAI(model="gemini-2.5-flash", temperature=0.2)
    
    contextualize_q_system_prompt = """Given a chat history and the latest user question \
    which might reference context in the chat history, formulate a standalone question \
    which can be understood without the chat history. Do NOT answer the question, \
    just reformulate it if needed and otherwise return it as is."""
    contextualize_q_prompt = ChatPromptTemplate.from_messages(
        [("system", contextualize_q_system_prompt), MessagesPlaceholder("chat_history"), ("human", "{input}")]
    )
    history_aware_retriever = create_history_aware_retriever(llm, retriever, contextualize_q_prompt)

    qa_system_prompt = """You are an expert financial analyst assistant. Your primary function is to answer questions based **only** on the financial document context provided to you.

    Your Rules:
    1.  Strictly Grounded: You must base your answers exclusively on the information found in the `Context` provided. Do not use any external knowledge or make assumptions.
    2.  Handle Missing Information: If the answer to a question cannot be found in the context, you **must** respond with: "The answer is not available in the provided document context."
    3.  Be Specific: When answering, if possible, refer to the specific financial statement (e.g., "According to the Profit & Loss statement...").
    4.  Quote Figures: If a question involves numbers or financial data, provide the exact figures and currency units (e.g., ₹ Crores) mentioned in the document.
    5.  Be Concise: Keep your answers direct and to the point.

    Context:
    {context}
    """
    
    qa_prompt = ChatPromptTemplate.from_messages(
        [
            ("system", qa_system_prompt),
            MessagesPlaceholder(variable_name="chat_history"),
            ("human", "{input}"),
        ]
    )
    question_answer_chain = create_stuff_documents_chain(llm, qa_prompt)

    rag_chain = create_retrieval_chain(history_aware_retriever, question_answer_chain)
    return rag_chain

# streamlit UI

if "analysis_report" not in st.session_state:
    st.session_state.analysis_report = ""
if "chat_history" not in st.session_state:
    st.session_state.chat_history = []
if "components_ready" not in st.session_state:
    st.session_state.components_ready = False

with st.sidebar:
    st.header("Setup")
    uploaded_file = st.file_uploader("Choose a DRHP PDF file", type="pdf")
    company_name_input = st.text_input("Enter the Company Name", placeholder="e.g. Reliance Industries Limited, SpaceX")

    if st.button("Process Document"):
        if uploaded_file is not None and company_name_input:
            with st.spinner("Analysis in progress... This may take a few minutes."):
                pdf_bytes = uploaded_file.read()

                llm_with_tools, tools, retriever = get_agent(pdf_bytes, company_name_input)
                st.session_state.llm_with_tools = llm_with_tools
                st.session_state.tools = tools
                st.session_state.company_name = company_name_input
                st.session_state.conversation_chain = get_conversational_chain(retriever)

                # reset chat history for new document
                st.session_state.chat_history = []
                st.session_state.analysis_report = ""

                st.success("Document processed! You can now use the analysis and chat features.")
                st.session_state.components_ready = True

        else:
            st.warning("Please upload a PDF and enter the company name.")


tab1, tab2 = st.tabs(["🤖 Full IPO Analysis", "💬 Chat with Document"])

with tab1:
    st.header("Automated IPO Analysis")
    st.write("After processing a document, click the button below to run the full analysis.")
    
    if st.button("Run Full Analysis"):
        if st.session_state.components_ready:
            with st.spinner("Agent is running the full analysis... This may take a few minutes."):
                tool_map = {tool.name: tool for tool in st.session_state.tools}
                master_query = f"""
                        Your mission is to generate a comprehensive IPO analysis report for **{company_name_input}**. Execute the following steps meticulously and synthesize the findings into a structured final report.

                        **Step 1: In-Depth Financial Statement Analysis**
                        Use the `financial_statement_search` tool to analyze the company's financials from the provided documents. You must investigate and report on the following:
                        -   **Balance Sheet:** Are reserves and surplus consistently growing? Are total borrowings consistently decreasing?
                        -   **Profit & Loss Statement:** Is the revenue from operations consistently growing? Is the profit after tax growing in proportion, or are margins shrinking?
                        -   **Cash Flow Statement:** Is the net cash flow from operating activities consistently positive?
                        
                        **IMPORTANT: If any of the financial statements are missing from the context, you must explicitly state that in your report and proceed with the analysis based on the available information.**

                        **Step 2: External Market and Valuation Analysis**
                        Use the `web_search` tool to gather the latest market data. You must find and report on:
                        -   The IPO's price band and the resulting P/E (Price-to-Earnings) ratio.
                        -   The average P/E ratio for the company's specific industry in India.
                        -   The P/E ratios of at least 2-3 key listed peer companies.
                        -   The current Grey Market Premium (GMP) for the IPO.

                        **Step 3: Synthesize and Generate the Final Report**
                        Once you have gathered all the necessary information, generate the final answer. The report **must** be structured exactly as follows:

                        ### **IPO Analysis: {company_name_input}**

                        **1. Financial Health Analysis:**
                        -   **Balance Sheet Verdict:** (Your findings on reserves and borrowings. Conclude: Strong, Average, or Weak).
                        -   **P&L Statement Verdict:** (Your findings on revenue and profit growth/margins. Conclude: Strong, Average, or Weak).
                        -   **Cash Flow Verdict:** (Your findings on cash from operations. Conclude: Positive, Negative, or **Data Not Available**).

                        **2. Valuation Analysis:**
                        -   **Valuation Verdict:** (Your comparison of the IPO's P/E against industry/peer P/Es. Conclude: Aggressive, Fairly Priced, or Attractive).

                        **3. Market Sentiment:**
                        -   **GMP Verdict:** (Your findings on the GMP and what it indicates).

                        **4. Key Strengths & Risks:**
                        -   **Strengths:** (A bulleted list of 2-3 key positive points).
                        -   **Risks:** (A bulleted list of 2-3 key risks or concerns).

                        **5. Final Verdict:**
                        -   (A clear recommendation: **SUBSCRIBE**, **AVOID**, or **SUBSCRIBE FOR LISTING GAINS**), followed by a concise justification.
                        """
                
                history = [HumanMessage(content=master_query)]
                final_answer = None
                for i in range(15):
                    response = st.session_state.llm_with_tools.invoke(history)
                    if not response.tool_calls:
                        final_answer = response.content
                        break
                    history.append(response)
                    for tool_call in response.tool_calls:
                        tool_name, tool_args = tool_call["name"], tool_call["args"]
                        if tool_name in tool_map:
                            observation = tool_map[tool_name].invoke(tool_args)
                        else:
                            observation = f"Error: Tool '{tool_name}' not found."
                        history.append(ToolMessage(content=str(observation), tool_call_id=tool_call["id"]))
                st.session_state.analysis_report = final_answer
                st.success("Analysis Complete!")
        else:
            st.warning("Please process a document in the sidebar first.")
    
    if st.session_state.analysis_report:
        st.markdown("---")
        st.subheader("Comprehensive IPO Analysis Report")
        st.markdown(st.session_state.analysis_report)

with tab2:
    st.header(f"Chat with the Document")
    st.write("Ask any question about the financial data in the uploaded PDF.")

    if not st.session_state.components_ready:
        st.warning("Please process a document in the sidebar first to enable chat.")
    else:
        for message in st.session_state.chat_history:
            with st.chat_message(message["role"]):
                st.markdown(message["content"])

        if user_question := st.chat_input("e.g., What was the total revenue in the last reported year?"):
            st.session_state.chat_history.append({"role": "user", "content": user_question})
            with st.chat_message("user"):
                st.markdown(user_question)

            with st.chat_message("assistant"):
                with st.spinner("Thinking..."):
                    response = st.session_state.conversation_chain.invoke(
                        {"chat_history": st.session_state.chat_history, "input": user_question}
                    )
                    bot_response = response["answer"]
                    st.markdown(bot_response)
            
            st.session_state.chat_history.append({"role": "assistant", "content": bot_response})