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

# Document categories
DOCUMENT_TYPES = ["Land Records", "Caste Certificates", "Property Registrations"]

# Set up session state
def initialize_session_state():
    if "processed_doc" not in st.session_state:
        st.session_state.processed_doc = None
    if "current_file" not in st.session_state:
        st.session_state.current_file = None
    if "gemini_api_key" not in st.session_state:
        st.session_state.gemini_api_key = None

# Reset session state
def reset_session_state():
    st.session_state.processed_doc = None
    st.session_state.current_file = None

# Encode file to base64
def encode_file(file):
    try:
        file_content = file.getvalue()
        return base64.b64encode(file_content).decode('utf-8')
    except Exception as e:
        st.error(f"Error encoding file: {str(e)}")
        return None

# Query Gemini API
def query_gemini(prompt, image_b64):
    if not st.session_state.gemini_api_key:
        st.error("API Key is missing. Please enter your Gemini API Key.")
        return None
    
    try:
        headers = {
            "Authorization": f"Bearer {st.session_state.gemini_api_key}",
            "Content-Type": "application/json"
        }
        payload = {
            "model": "gemini-1.5-pro",  # Change based on the latest API model
            "prompt": prompt,
            "image": image_b64
        }
        response = requests.post(
            "https://generativelanguage.googleapis.com/v1/models/gemini-1.5-pro:generateContent",  # Corrected URL
            headers=headers,
            json=payload
        )
        response.raise_for_status()
        return response.json().get("result", "No response received.")
    except Exception as e:
        st.error(f"Gemini API error: {str(e)}")
        return None

# Process the document
def process_document(file):
    try:
        with st.spinner("Analyzing document..."):
            image_b64 = encode_file(file)
            if not image_b64:
                return

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

            # 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 and organize all important details from this document including:
            - Names
            - Dates
            - Identification numbers
            - Locations
            - Key terms
            Format as a bullet-point list with clear headings."""
            details = query_gemini(extract_prompt, image_b64)

            # Verify authenticity
            verify_prompt = "Analyze this document for signs of tampering or forgery. Provide verification results."
            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",
                "preview": st.session_state.doc_preview
            }

            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 application
def main():
    st.set_page_config(page_title="DocVerify AI", layout="wide")
    initialize_session_state()
    
    # Sidebar Controls
    with st.sidebar:
        st.header("Document Controls")

        # API Key input
        st.session_state.gemini_api_key = st.text_input("Enter Gemini API Key", type="password")

        uploaded_file = st.file_uploader(
            "Upload Document",
            type=["pdf", "jpg", "jpeg", "png"],
            key="uploaded_file"
        )
        
        if st.button("Process Document"):
            if uploaded_file:
                st.session_state.current_file = uploaded_file
                process_document(uploaded_file)
            else:
                st.error("Please select a document to process.")
        
        if st.button("New Document"):
            reset_session_state()
            st.rerun()
        
        if st.session_state.processed_doc:
            st.divider()
            st.subheader("Document Summary")
            st.markdown(f"**Type:** {st.session_state.processed_doc['type']}")
            st.markdown(f"**Verification Status:**\n{st.session_state.processed_doc['verification']}")

    # Main Interface
    st.title("📄 Automated Document Verifier")
    
    if st.session_state.processed_doc and 'preview' in st.session_state.processed_doc:
        col1, col2 = st.columns([1, 2])
        with col1:
            st.subheader("Document Preview")
            st.image(st.session_state.processed_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 and click 'Process Document' to start verification.")

if __name__ == "__main__":
    main()