import streamlit as st
import pdfplumber
import pytesseract
import pandas as pd
from PIL import Image
import cv2
import numpy as np
import os
import json
from groq import Groq

st.set_page_config(page_title="AI Invoice Intelligence Engine", layout="wide")

st.title("🤖 AI Financial Document Intelligence")
st.write("Upload ANY invoice / receipt / subscription / purchase document")

uploaded_file = st.file_uploader(
    "Upload File",
    type=["pdf", "jpg", "jpeg", "png"]
)

# ---------------------------------------------------------
# TEXT EXTRACTION
# ---------------------------------------------------------

def extract_text(uploaded_file):

    text = ""

    if uploaded_file.type == "application/pdf":
        with pdfplumber.open(uploaded_file) as pdf:
            for page in pdf.pages:
                page_text = page.extract_text()
                if page_text:
                    text += page_text + "\n"
    else:
        image = Image.open(uploaded_file)
        img_np = np.array(image)
        gray = cv2.cvtColor(img_np, cv2.COLOR_BGR2GRAY)
        text = pytesseract.image_to_string(gray)

    return text


# ---------------------------------------------------------
# SAFE JSON LOADER
# ---------------------------------------------------------

def safe_json_load(text):

    text = text.strip()
    text = text.replace("```json", "")
    text = text.replace("```", "")

    try:
        return json.loads(text)
    except:
        return None


# ---------------------------------------------------------
# SPLIT LARGE DOCUMENT INTO CHUNKS
# ---------------------------------------------------------

def split_into_chunks(text, max_chars=4000):

    chunks = []
    current = ""

    for line in text.split("\n"):
        if len(current) + len(line) > max_chars:
            chunks.append(current)
            current = ""
        current += line + "\n"

    if current:
        chunks.append(current)

    return chunks


# ---------------------------------------------------------
# AI EXTRACTION (GROQ)
# ---------------------------------------------------------

def ai_extract_invoice(text):

    client = Groq(api_key=os.getenv("GROQ_API_KEY"))

    chunks = split_into_chunks(text)

    invoices = []

    for chunk in chunks:

        prompt = f"""
        You are a financial document extraction engine.

        If this text contains:
        - invoice
        - receipt
        - subscription
        - payment confirmation
        - order
        - tax invoice

        Extract structured JSON array.

        Format:

        [
          {{
            "document_type": "",
            "invoice_or_order_number": "",
            "invoice_date": "",
            "vendor_name": "",
            "vendor_gstin": null,
            "taxable_amount": null,
            "cgst": null,
            "sgst": null,
            "igst": null,
            "total_amount": null
          }}
        ]

        If no financial transaction, return [].

        TEXT:
        {chunk}
        """

        try:
            completion = client.chat.completions.create(
                model="llama-3.1-8b-instant",
                messages=[{"role": "user", "content": prompt}],
                temperature=0
            )

            response_text = completion.choices[0].message.content
            parsed = safe_json_load(response_text)

            if parsed and isinstance(parsed, list):
                invoices.extend(parsed)

        except Exception as e:
            print("Chunk error:", e)
            continue

    return invoices


# ---------------------------------------------------------
# CONFIDENCE SCORE
# ---------------------------------------------------------

def calculate_confidence(row):

    total_fields = len(row)
    filled_fields = sum(
        1 for v in row.values
        if v not in [None, "", "null"]
    )

    return round((filled_fields / total_fields) * 100, 2)


# ---------------------------------------------------------
# MAIN EXECUTION
# ---------------------------------------------------------

if uploaded_file:

    raw_text = extract_text(uploaded_file)

    st.subheader("📄 Raw Text Preview")
    st.text_area("Preview", raw_text[:3000], height=250)

    with st.spinner("🤖 AI Analyzing Document..."):
        invoices = ai_extract_invoice(raw_text)

    if invoices:

        df = pd.DataFrame(invoices).drop_duplicates()
        df["confidence_%"] = df.apply(calculate_confidence, axis=1)

        st.success(f"✅ {len(df)} document(s) detected")
        st.dataframe(df, use_container_width=True)

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

    else:
        st.warning("⚠ No financial transaction detected.")