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

# Configuration
GEMINI_API_KEY = st.secrets["GEMINI_API_KEY"]
GEMINI_MODEL = "gemini-1.5-flash"  # Verify with Gemini API documentation
DOCUMENT_TYPES = ["Land Records", "Caste Certificates", "Property Registrations"]

# Define a class to mimic Streamlit's UploadedFile object for the example image
class ExampleFile:
    def __init__(self, file_path):
        self.name = os.path.basename(file_path)
        self.type = "image/jpeg"
        with open(file_path, "rb") as f:
            self._buffer = BytesIO(f.read())
    
    def getvalue(self):
        return self._buffer.getvalue()

# Initialize 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

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

# Encode file to base64 for Gemini API
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):
    try:
        headers = {
            "Authorization": f"Bearer {GEMINI_API_KEY}",
            "Content-Type": "application/json"
        }
        payload = {
            "model": GEMINI_MODEL,
            "prompt": prompt,
            "image": image_b64
        }
        response = requests.post(
            "https://api.gemini.com/v1/analyze",  # Replace with actual Gemini API endpoint
            headers=headers,
            json=payload
        )
        response.raise_for_status()
        return response.json().get("result", "")
    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..."):
            # Encode file to base64
            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()))
            elif file.type == "text/plain":
                text = file.getvalue().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))
                st.session_state.doc_preview = img
            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. Check for: inconsistent fonts, alignment issues, suspicious modifications. Provide verification conclusion."
            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")
        
        uploaded_file = st.file_uploader(
            "Upload Document",
            type=["pdf", "jpg", "jpeg", "png", "txt"],
            key="uploaded_file"
        )
        
        if st.button("Use Example Image"):
            example_file_path = "Caste-Certificate.jpg"
            st.session_state.current_file = ExampleFile(example_file_path)
        elif uploaded_file:
            st.session_state.current_file = uploaded_file
        else:
            st.session_state.current_file = None
        
        if st.button("Process Document"):
            if st.session_state.current_file:
                process_document(st.session_state.current_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 select a document and click 'Process Document' to begin verification analysis.")

if __name__ == "__main__":
    main()