app.py

import re
import json
import os
import time
from datetime import datetime, date
from io import BytesIO
import pandas as pd
import streamlit as st
import google.generativeai as genai
import pypdf
from fpdf import FPDF
from google.api_core import exceptions
import markdown
from bs4 import BeautifulSoup

# Configure API key for Gemini - Ensure this is set in your environment variables
api_key = os.getenv('Gemini')

def configure_gemini(api_key):
    """
    Configures the Gemini model for transaction extraction as specified by the user.
    """
    st.info("Configuring Gemini API for transaction extraction...")
    genai.configure(api_key=api_key)
    return genai.GenerativeModel('gemini-2.0-flash-thinking-exp')

def configure_gemini1(api_key):
    """
    Configures the Gemini model for report generation as specified by the user.
    """
    st.info("Configuring Gemini API for report generation...")
    genai.configure(api_key=api_key)
    return genai.GenerativeModel('gemini-2.5-pro')

def read_pdf_pages(file_obj):
    st.info(f"Reading PDF pages from {file_obj.name}...")
    file_obj.seek(0)
    pdf_reader = pypdf.PdfReader(file_obj)
    total_pages = len(pdf_reader.pages)
    st.info(f"Found {total_pages} pages in PDF.")
    return pdf_reader, total_pages

def extract_page_text(pdf_reader, page_num):
    if page_num < len(pdf_reader.pages):
        text = pdf_reader.pages[page_num].extract_text()
        if not text or not text.strip():
            st.warning(f"Page {page_num + 1} appears to be empty or contains no extractable text.")
        return text if text else ""
    return ""

def process_with_gemini(model, text):
    prompt = """Analyze this bank statement and extract transactions in JSON format with these fields:
    - Date (format DD/MM/YYYY)
    - Description
    - Amount (just the integer value)
    - Type (is 'income' if 'credit amount', else 'expense')
    - Customer Name (Only If Type is 'income' and if no name is extracted write 'general income' and if type is not 'income' write 'expense')
    - City (In address of bank statement)
    - Category_of_expense (a string, if transaction 'Type' is 'expense' categorize it based on description into: Water and electricity, Salaries and wages, Repairs & Maintenance, Motor vehicle expenses, Projects Expenses, Hardware expenses, Refunds, Accounting fees, Loan interest, Bank charges, Insurance, SARS PAYE UIF, Advertising & Marketing, Logistics and distribution, Fuel, Website hosting fees, Rentals, Subscriptions, Computer internet and Telephone, Staff training, Travel and accommodation, Depreciation, Other expenses. If no category matches, default to 'Other expenses'. If 'Type' is 'income' set Destination_of_funds to 'income'.)
    - ignore opening or closing balances, charts and analysis.

    Return ONLY valid JSON with this structure:
    {
        "transactions": [
            {
                "Date": "string",
                "Description": "string",
                "Customer_name": "string",
                "City": "string",
                "Amount": number,
                "Type": "string",
                "Category_of_expense": "string"
            }
        ]
    }"""
    try:
        response = model.generate_content([prompt, text])
        time.sleep(6)
        return response.text
    except exceptions.GoogleAPICallError as e:
        st.error(f"A Google API call error occurred during transaction extraction: {e}")
        return None
    except Exception as e:
        st.error(f"An unexpected error occurred during Gemini transaction extraction: {e}")
        return None

def process_pdf_pages(model, pdf_reader, total_pages, progress_callback=None):
    all_transactions = []
    for page_num in range(total_pages):
        if progress_callback:
            progress_callback(page_num / total_pages, f"Processing page {page_num + 1} of {total_pages}")
        page_text = extract_page_text(pdf_reader, page_num)
        if not page_text.strip():
            continue
        json_response = process_with_gemini(model, page_text)
        if json_response:
            match = re.search(r'\{.*\}', json_response, re.DOTALL)
            if not match:
                continue
            json_str = match.group(0)
            try:
                data = json.loads(json_str)
                transactions = data.get('transactions', [])
                if transactions:
                    all_transactions.extend(transactions)
            except json.JSONDecodeError:
                continue
    return all_transactions

def aggregate_financial_data(transactions: list, statement_type: str):
    st.info(f"Performing local financial aggregation for {len(transactions)} transactions...")
    if not transactions:
        return None
    df = pd.DataFrame(transactions)
    if 'Amount' not in df.columns:
        return None
    df['Amount'] = df['Amount'].astype(str).str.replace(r'[^\d.]', '', regex=True)
    df['Amount'] = pd.to_numeric(df['Amount'], errors='coerce').fillna(0)
    df['Type'] = df['Type'].str.lower()
    total_income = df[df['Type'] == 'income']['Amount'].sum()
    total_expenses = df[df['Type'] == 'expense']['Amount'].sum()
    net_position = total_income - total_expenses
    aggregated_data = {
        "total_income": round(total_income, 2),
        "total_expenses": round(total_expenses, 2),
        "net_position": round(net_position, 2),
        "transaction_count": len(df)
    }
    if statement_type == "Income Statement":
        aggregated_data["expense_breakdown"] = df[df['Type'] == 'expense'].groupby('Category_of_expense')['Amount'].sum().round(2).to_dict()
        aggregated_data["income_breakdown"] = df[df['Type'] == 'income'].groupby('Customer_name')['Amount'].sum().round(2).to_dict()
    st.success("Local financial aggregation complete.")
    return aggregated_data

def generate_financial_report(model, aggregated_data, start_date, end_date, statement_type):
    """
    Generates a financial report using a simplified, high-level prompt that
    trusts the model to create the correct structure and avoids using any
    Markdown characters that could break rendering.
    """
    st.info(f"Preparing to generate {statement_type} with pre-aggregated data...")

    prompt = f"""You are an expert financial analyst. Your task is to generate a professional Income Statement in Markdown format using the pre-aggregated JSON data provided below.

Here is the financial data:
{json.dumps(aggregated_data, indent=2)}

Your instructions for the report are:
The main title of the report is "Income Statement".
The reporting period is from {start_date.strftime('%d %B %Y')} to {end_date.strftime('%d %B %Y')}.
The currency is South African Rand (ZAR).
The report must contain sections for Revenue, Operating Expenses, and Net Income or Loss. Each of these sections must be a clear table.
The Revenue section must contain a single table showing only the 'Total Revenue'. Do NOT create a table that itemizes or lists individual income sources.
The report must also include a "Key Highlights" section with bullet points and a final "Summary" paragraph.
Use the provided JSON data for all financial figures.
For the Net Income or Loss table, if the net position is negative, display the amount in parentheses.
Separate the major sections with a horizontal rule.
"""
    try:
        st.info("Sending request to Gemini for final report formatting...")
        response = model.generate_content([prompt])
        st.success("Successfully received formatted financial report from Gemini.")
        return response.text
    except exceptions.GoogleAPICallError as e:
        st.error(f"A Google API call error occurred during report generation: {e}")
        return None
    except Exception as e:
        st.error(f"An unexpected error occurred during Gemini report generation: {e}")
        return None

def create_pdf_report(report_text):
    """
    Creates a PDF from markdown text. This version includes fixes for both the
    ValueError from incorrect int conversion and the table rendering logic.
    """
    if not report_text:
        st.warning("Report text is empty, skipping PDF generation.")
        raise ValueError("Input report_text cannot be empty.")
    try:
        st.info("Starting PDF generation from markdown report...")
        cleaned_md = re.sub(r'```markdown|```', '', report_text, flags=re.MULTILINE).strip()
        html_content = markdown.markdown(cleaned_md, extensions=['tables'])
        soup = BeautifulSoup(html_content, 'html.parser')
        
        pdf = FPDF()
        pdf.set_auto_page_break(auto=True, margin=15)
        pdf.set_left_margin(15)
        pdf.set_right_margin(15)
        pdf.add_page()
        
        for element in soup.find_all(True):
            if element.name in ['h1', 'h2', 'h3']:
                level = int(element.name[1])
                font_size = {1: 16, 2: 14, 3: 12}.get(level, 11)
                pdf.set_font('helvetica', 'B', font_size)
                pdf.multi_cell(0, 10, element.get_text().strip())
                pdf.ln(level * 2)
            elif element.name == 'p':
                pdf.set_font('helvetica', '', 11)
                pdf.multi_cell(0, 6, element.get_text().strip())
                pdf.ln(4)
            elif element.name == 'i':
                 pdf.set_font('helvetica', 'I', 11)
                 pdf.multi_cell(0, 6, element.get_text().strip())
                 pdf.ln(4)
            elif element.name == 'hr':
                pdf.line(pdf.get_x(), pdf.get_y(), pdf.w - pdf.r_margin, pdf.get_y())
                pdf.ln(5)
            elif element.name == 'ul':
                pdf.ln(2)
                for li in element.find_all('li'):
                    pdf.set_font('helvetica', '', 11)
                    item_text = li.get_text().strip().replace('•', '-')
                    pdf.multi_cell(0, 5, f"  -  {item_text}")
                    pdf.ln(1)
                pdf.ln(4)
            elif element.name == 'table':
                header = [th.get_text().strip() for th in element.find_all('th')]
                rows = [[td.get_text().strip() for td in tr.find_all('td')] for tr in element.find_all('tr')[1:]]
                
                if header:
                    pdf.set_font('helvetica', 'B', 10)
                    pdf.set_fill_color(230, 230, 230)
                    col_widths = [ (pdf.w - pdf.l_margin - pdf.r_margin) * 0.6, (pdf.w - pdf.l_margin - pdf.r_margin) * 0.4 ]
                    for i, header_text in enumerate(header):
                        pdf.cell(col_widths[i], 8, header_text, border=1, fill=True, align='C')
                    pdf.ln()

                pdf.set_font('helvetica', '', 10)
                for row in rows:
                    is_total_row = any('Total' in cell for cell in row)
                    if is_total_row:
                        pdf.set_font('helvetica', 'B', 10)

                    if len(row) == len(col_widths):
                        pdf.cell(col_widths[0], 7, row[0], border=1)
                        pdf.cell(col_widths[1], 7, row[1], border=1, align='R')
                        pdf.ln()
                    
                    if is_total_row:
                        pdf.set_font('helvetica', '', 10)
                pdf.ln(6)

        st.info("Content added to PDF. Outputting PDF to buffer...")
        
        pdf_output = pdf.output()

        st.success("PDF report generated successfully.")
        return BytesIO(pdf_output)
    except Exception as e:
        st.error(f"Failed to generate PDF: {e}")
        st.exception(e)
        raise

def main():
    st.title("Quantitlytix AI")
    st.markdown("*Bank Statement Parser & Financial Report Generator*")

    if 'min_date' not in st.session_state:
        st.session_state['min_date'] = date(2024, 1, 1)
    if 'max_date' not in st.session_state:
        st.session_state['max_date'] = date.today()
    if 'transactions' not in st.session_state:
        st.session_state['transactions'] = []

    input_type = st.sidebar.radio("Select Input Type", ("Bulk Bank Statement Upload", "CSV Upload"))

    if input_type == "Bulk Bank Statement Upload":
        uploaded_files = st.file_uploader("Upload PDF bank statements", type="pdf", accept_multiple_files=True)
        if uploaded_files:
            model = configure_gemini(api_key)
            progress_bar = st.progress(0)
            all_transactions = []
            for i, file in enumerate(uploaded_files):
                st.text(f"Processing {file.name}...")
                pdf_reader, total_pages = read_pdf_pages(file)
                if total_pages > 0:
                    file_transactions = process_pdf_pages(model, pdf_reader, total_pages)
                    all_transactions.extend(file_transactions)
                progress_bar.progress((i + 1) / len(uploaded_files))
            st.session_state['transactions'] = all_transactions
            st.success(f"All PDF files processed. Total transactions collected: {len(st.session_state['transactions'])}.")

    elif input_type == "CSV Upload":
        uploaded_csv = st.file_uploader("Upload CSV of transactions", type="csv")
        if uploaded_csv:
            df = pd.read_csv(uploaded_csv)
            df = df.loc[:, ~df.columns.str.startswith('Unnamed:')]
            st.session_state['transactions'] = df.to_dict(orient='records')
            st.success(f"Successfully loaded {len(st.session_state['transactions'])} transactions from CSV.")

    if st.session_state['transactions']:
        df = pd.DataFrame(st.session_state['transactions'])
        # Clean and validate data
        df['Date'] = pd.to_datetime(df['Date'], errors='coerce', dayfirst=True)
        df.dropna(subset=['Date'], inplace=True)
        
        # Perform the same robust cleaning for Amount as in the aggregation function
        if 'Amount' in df.columns:
            df['Amount'] = df['Amount'].astype(str).str.replace(r'[^\d.]', '', regex=True)
            df['Amount'] = pd.to_numeric(df['Amount'], errors='coerce').fillna(0)

        if not df.empty:
            st.session_state['min_date'] = df['Date'].min().date()
            st.session_state['max_date'] = df['Date'].max().date()
        
        st.write("### Extracted Transactions")
        st.dataframe(df.astype(str))

        # --- NEW FEATURE: Download Processed CSV Button ---
        st.download_button(
            label="Download Processed CSV",
            data=df.to_csv(index=False).encode('utf-8'),
            file_name='processed_transactions.csv',
            mime='text/csv',
        )
        # --- END OF NEW FEATURE ---

    st.write("### Generate Financial Report")
    col1, col2 = st.columns(2)
    with col1:
        start_date = st.date_input("Start Date", st.session_state['min_date'])
    with col2:
        end_date = st.date_input("End Date", st.session_state['max_date'])
    statement_type = st.selectbox("Select Financial Statement", ["Income Statement"])

    if st.button("Generate Financial Report"):
        if not st.session_state['transactions']:
            st.error("No transactions available to generate report. Please upload files first.")
        else:
            # Re-create dataframe from session state to ensure it's fresh for filtering
            df_report = pd.DataFrame(st.session_state['transactions'])
            df_report['Date'] = pd.to_datetime(df_report['Date'], errors='coerce', dayfirst=True)
            
            mask = (df_report['Date'] >= pd.to_datetime(start_date)) & (df_report['Date'] <= pd.to_datetime(end_date))
            filtered_df = df_report.loc[mask]

            if filtered_df.empty:
                st.warning("No transactions found within the selected date range.")
            else:
                st.info(f"Found {len(filtered_df)} transactions within the selected date range.")
                filtered_transactions_list = filtered_df.to_dict(orient='records')
                try:
                    with st.spinner("Aggregating financial data locally..."):
                        aggregated_summary = aggregate_financial_data(filtered_transactions_list, statement_type)
                    if aggregated_summary:
                        with st.spinner("Generating formatted report with Gemini..."):
                            model1 = configure_gemini1(api_key)
                            report_text = generate_financial_report(model1, aggregated_summary, start_date, end_date, statement_type)
                        if report_text:
                            st.success("Financial report generated successfully!")
                            st.markdown("### Financial Report Preview")
                            st.markdown(report_text, unsafe_allow_html=True)
                            pdf_buffer = create_pdf_report(report_text)
                            st.download_button(
                                label="Download Financial Report as PDF",
                                data=pdf_buffer,
                                file_name=f"{statement_type.replace(' ', '_')}_{datetime.now().strftime('%Y%m%d')}.pdf",
                                mime="application/pdf"
                            )
                except Exception as e:
                    st.error(f"An unexpected error occurred during the report generation process: {e}")
                    st.exception(e)

if __name__ == "__main__":
    main()