import streamlit as st
from PIL import Image
import pytesseract
import pandas as pd
import re
import cv2
import numpy as np
from PyPDF2 import PdfReader
import tempfile
import os

st.set_page_config(page_title="CA Invoice Automation", layout="wide")

st.title("📑 CA Invoice Automation App")

uploaded_file = st.file_uploader(
    "Upload Invoice (Image or PDF)",
    type=["jpg", "jpeg", "png", "pdf"]
)

# --------------------------------
# Classification
# --------------------------------
def classify_invoice(text):
    if "tax invoice" in text.lower():
        return "GST_PURCHASE"
    elif "credit note" in text.lower():
        return "CREDIT_NOTE"
    else:
        return "UNKNOWN"

# --------------------------------
# Extraction
# --------------------------------
def extract_data(text):
    gstin_pattern = r"\d{2}[A-Z]{5}\d{4}[A-Z]\dZ\d"
    date_pattern = r"\d{2}[-/]\d{2}[-/]\d{4}"
    amount_pattern = r"\d{1,3}(,\d{3})*(\.\d{2})?"

    gstin = re.search(gstin_pattern, text)
    date = re.search(date_pattern, text)
    amount = re.search(amount_pattern, text)

    return {
        "Date": date.group() if date else None,
        "GSTIN": gstin.group() if gstin else None,
        "Total Amount": amount.group() if amount else None
    }

# --------------------------------
# OCR for Images
# --------------------------------
def ocr_image(pil_image):
    image_np = np.array(pil_image)
    gray = cv2.cvtColor(image_np, cv2.COLOR_BGR2GRAY)
    return pytesseract.image_to_string(gray)

# --------------------------------
# MAIN
# --------------------------------
if uploaded_file:

    full_text = ""

    try:

        # ================================
        # -------- PDF HANDLING ----------
        # ================================
        if uploaded_file.type == "application/pdf":

            st.info("Extracting text directly from PDF...")

            with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp:
                tmp.write(uploaded_file.getvalue())
                tmp_path = tmp.name

            reader = PdfReader(tmp_path)

            for page in reader.pages[:5]:  # Limit to first 5 pages
                extracted_text = page.extract_text()
                if extracted_text:
                    full_text += extracted_text + "\n"

            os.remove(tmp_path)

            if not full_text.strip():
                st.warning("No embedded text found in PDF (might be scanned).")

        # ================================
        # -------- IMAGE HANDLING --------
        # ================================
        else:
            image = Image.open(uploaded_file)
            st.image(image, caption="Uploaded Invoice", use_column_width=True)
            full_text = ocr_image(image)

        # --------------------------------
        # Show OCR / Extracted Text
        # --------------------------------
        st.subheader("📝 Extracted Text")
        st.text(full_text)

        # --------------------------------
        # Classification
        # --------------------------------
        invoice_type = classify_invoice(full_text)
        st.success(f"Invoice Type: {invoice_type}")

        # --------------------------------
        # Extraction
        # --------------------------------
        data = extract_data(full_text)

        formatted_data = {
            "Date": data["Date"],
            "Voucher Type": "Purchase" if invoice_type=="GST_PURCHASE" else "Unknown",
            "Vendor": "To be detected",
            "GSTIN": data["GSTIN"],
            "Invoice No": "To be extracted",
            "Taxable": None,
            "CGST": None,
            "SGST": None,
            "Total": data["Total Amount"],
            "Debit Ledger": "Purchase A/c",
            "Credit Ledger": "Vendor A/c",
            "Narration": f"Being purchase invoice dated {data['Date']}",
        }

        df = pd.DataFrame([formatted_data])

        st.subheader("📊 Structured Voucher Data")
        st.dataframe(df, use_container_width=True)

        st.download_button(
            "⬇ Download CSV",
            df.to_csv(index=False),
            "voucher_output.csv",
            "text/csv"
        )

    except Exception as e:
        st.error("Processing failed.")
        st.write(str(e))