app.py

import streamlit as st
import base64
import requests
from PIL import Image, ImageDraw
from io import BytesIO
import fitz  # PyMuPDF
import time

# Configuration - Get API key from Streamlit secrets
GEMINI_API_KEY = st.secrets["GEMINI_API_KEY"]
GEMINI_MODEL = "gemini-2.0-flash"
DOCUMENT_TYPES = ["Land Records", "Caste Certificates", "Property Registrations", "Others"]

# Initialize session state (excluding widget-controlled keys)
def initialize_session_state():
    if "chat_history" not in st.session_state:
        st.session_state["chat_history"] = []
    if "processed_doc" not in st.session_state:
        st.session_state["processed_doc"] = None
    if "doc_preview" not in st.session_state:
        st.session_state["doc_preview"] = None

# Reset session state (excluding widget-controlled keys)
def reset_session_state():
    for key in ["chat_history", "processed_doc", "doc_preview"]:
        st.session_state.pop(key, None)

# Encode uploaded file to base64
def encode_file(uploaded_file):
    try:
        file_bytes = uploaded_file.getvalue()
        
        if uploaded_file.type == "application/pdf":
            pdf = fitz.open(stream=BytesIO(file_bytes))
            page = pdf[0]
            pix = page.get_pixmap()
            img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
        elif uploaded_file.type.startswith("image/"):
            img = Image.open(BytesIO(file_bytes))
        elif uploaded_file.type == "text/plain":
            text = file_bytes.decode("utf-8")
            img = Image.new("RGB", (800, 600), color=(73, 109, 137))
            d = ImageDraw.Draw(img)
            d.text((10, 10), text, fill=(255, 255, 0))
        else:
            st.error("Unsupported file format")
            return None

        img_byte_arr = BytesIO()
        img.save(img_byte_arr, format="JPEG")
        return base64.b64encode(img_byte_arr.getvalue()).decode("utf-8")
    
    except Exception as e:
        st.error(f"File processing error: {str(e)}")
        return None

# Query Gemini API
def query_gemini(prompt, image_b64=None):
    url = f"https://generativelanguage.googleapis.com/v1/models/{GEMINI_MODEL}:generateContent?key={GEMINI_API_KEY}"
    
    parts = [{"text": prompt}]
    if image_b64:
        parts.append({
            "inline_data": {
                "mime_type": "image/jpeg",
                "data": image_b64
            }
        })
    
    try:
        response = requests.post(
            url,
            json={"contents": [{"parts": parts}]},
            headers={"Content-Type": "application/json"},
            timeout=30
        )
        if response.status_code != 200:
            st.error(f"API Request failed with status code: {response.status_code}")
            return None
        
        data = response.json()
        if "error" in data:
            st.error(f"API Error: {data['error'].get('message', 'Unknown error')}")
            return None
            
        if not data.get("candidates"):
            st.error("No response candidates found in API response")
            return None
            
        candidate = data["candidates"][0]
        return candidate.get("content", {}).get("parts", [{}])[0].get("text", "No response text found")
    
    except requests.exceptions.RequestException as e:
        st.error(f"API Request failed: {str(e)}")
        return None
    except Exception as e:
        st.error(f"Unexpected error: {str(e)}")
        return None

# Process the uploaded document
def process_document():
    uploaded_file = st.session_state.get("uploaded_file")
    if not uploaded_file:
        st.error("Please upload a document first.")
        return
    
    try:
        with st.spinner("Analyzing document..."):
            # Encode file to base64
            image_b64 = encode_file(uploaded_file)
            if not image_b64:
                return

            # Store preview image
            if uploaded_file.type == "application/pdf":
                pdf = fitz.open(stream=BytesIO(uploaded_file.getvalue()))
                page = pdf[0]
                pix = page.get_pixmap()
                st.session_state.doc_preview = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
            else:
                st.session_state.doc_preview = Image.open(uploaded_file)

            # Classify document
            classify_prompt = (
                f"Classify this document into one of these categories: {', '.join(DOCUMENT_TYPES)}. "
                "Respond only with the category name."
            )
            doc_type = query_gemini(classify_prompt, image_b64)
            
            # Extract details
            extract_prompt = (
                "Extract key details including:\n"
                "- Names\n"
                "- Dates\n"
                "- Identification numbers\n"
                "- Locations\n"
                "Format as a bullet-point list."
            )
            details = query_gemini(extract_prompt, image_b64)

            # Verify authenticity
            verify_prompt = "Analyze this document for signs of tampering. Provide verification status in short(2 Lines)."
            verification = query_gemini(verify_prompt, image_b64)

            st.session_state.processed_doc = {
                "type": doc_type or "Unclassified",
                "details": details or "No details extracted",
                "verification": verification or "Verification failed",
            }

            st.success("Document processing complete!")
            time.sleep(1)
    
    except Exception as e:
        st.error(f"Document processing failed: {str(e)}")
        st.session_state.processed_doc = None

# Main app function
def main():
    st.set_page_config(page_title="DocVerify AI", layout="wide")
    initialize_session_state()
    
    # Sidebar Controls
    with st.sidebar:
        st.header("Document Controls")
        # The file uploader widget manages its own state with key "uploaded_file"
        st.file_uploader(
            "Upload Document",
            type=["pdf", "jpg", "jpeg", "png", "txt"],
            key="uploaded_file",
            on_change=process_document,
            help="Supported formats: PDF, JPG, PNG, TXT"
        )
        
        if st.button("New Document"):
            reset_session_state()
            st.experimental_rerun()
        
        if st.session_state.get("processed_doc"):
            st.divider()
            st.subheader("Document Summary")
            st.markdown(f"**Type:** {st.session_state.processed_doc['type']}")
            st.markdown(f"**Verification Status:** {st.session_state.processed_doc['verification']}")

    # Main Interface
    st.title("📄 Automated Document Verifier")
    
    if st.session_state.get("processed_doc") and st.session_state.get("doc_preview"):
        col1, col2 = st.columns([1, 2])
        with col1:
            st.subheader("Document Preview")
            st.image(st.session_state.doc_preview, use_column_width=True)
        
        with col2:
            st.subheader("Extracted Details")
            st.markdown(st.session_state.processed_doc["details"])
            
            st.subheader("Verification Analysis")
            st.markdown(st.session_state.processed_doc["verification"])
    else:
        st.info("Please upload a document to begin verification analysis")

if __name__ == "__main__":
    main()