import base64
import streamlit as st
from openai import OpenAI
import pdf2image
from PIL import Image
import io
import tempfile

# Load environment variables

# Initialize OpenAI client
client = OpenAI()



def convert_pdf_to_images(pdf_file):
    """Convert PDF to list of images"""
    with tempfile.NamedTemporaryFile(delete=False, suffix='.pdf') as tmp_file:
        tmp_file.write(pdf_file.getvalue())
        pdf_path = tmp_file.name
    
    images = pdf2image.convert_from_path(pdf_path)
    os.unlink(pdf_path)
    return images

def format_response(text):
    """Format the analysis response with clean styling"""
    formatted_text = """
    <div style="
        background-color: white;
        padding: 20px;
        border-radius: 5px;
        font-family: Arial, sans-serif;
        box-shadow: 0 2px 4px rgba(0,0,0,0.1);
    ">
    """
    
    # Split into pages
    pages = text.split("Page")
    
    for page_num, page_content in enumerate(pages[1:], 1):  # Skip first empty split
        # Add page header
        formatted_text += f'<div style="margin-bottom: 30px;">'
        formatted_text += f'<h3 style="color: #2c3e50; margin-bottom: 15px;">Page {page_num}</h3>'
        
        # Process each line
        lines = page_content.split('\n')
        for line in lines:
            # Skip empty lines and lines with asterisks
            if line.strip() and not line.strip().startswith('*') and not line.strip().startswith('Here'):
                # Remove asterisks and dashes
                line = line.replace('**', '').replace('- ', '')
                
                if ':' in line:
                    label, value = line.split(':', 1)
                    formatted_text += f'<div style="margin-bottom: 10px; display: flex;">'
                    formatted_text += f'<span style="font-weight: bold; color: #2c3e50; min-width: 200px;">{label.strip()}</span>'
                    formatted_text += f'<span style="color: #34495e; flex: 1;">{value.strip()}</span>'
                    formatted_text += '</div>'
        
        formatted_text += '</div>'
        
        # Add separator between pages except for the last page
        if page_num < len(pages) - 1:
            formatted_text += '<hr style="border: 1px solid #eee; margin: 20px 0;">'
    
    formatted_text += '</div>'
    return formatted_text

def analyze_image(image):
    """Analyze image using OpenAI API"""
    try:
        img_byte_arr = io.BytesIO()
        image.save(img_byte_arr, format='PNG')
        img_byte_arr = img_byte_arr.getvalue()
        
        base64_image = base64.b64encode(img_byte_arr).decode("utf-8")
        
        response = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "text",
                            "text": """Please analyze the image and extract the following information:
                            - Sender information
                            - Recipient information
                            - Container details
                            - Weights and measurements
                            - Dates and reference numbers
                            - Cargo details
                            
                            Format the response as 'Label: Value' pairs."""
                        },
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:image/jpeg;base64,{base64_image}"
                            },
                        },
                    ],
                }
            ],
            max_tokens=1000
        )
        
        return response.choices[0].message.content
    except Exception as e:
        return f"An error occurred: {str(e)}"

def main():
    st.set_page_config(page_title="Document Analysis App", layout="wide")
    
    # Custom CSS to set light background and improve button styling
    st.markdown("""
        <style>
        .stApp {
            background-color: white;
        }
        .stButton>button {
            width: 100%;
            background-color: #2c3e50;
            color: white;
            border: none;
            padding: 10px 20px;
            border-radius: 5px;
            margin-top: 20px;
            margin-bottom: 20px;
        }
        .stButton>button:hover {
            background-color: #34495e;
        }
        .uploadedFile {
            margin-bottom: 20px;
        }
        </style>
    """, unsafe_allow_html=True)
    
    col1, col2, col3 = st.columns([1,2,1])
    with col2:
        st.title("Document Analysis App")
        
        uploaded_file = st.file_uploader("Upload document (PDF/Image)", type=['pdf', 'png', 'jpg', 'jpeg'])
        
        if uploaded_file is not None:
            if uploaded_file.type == "application/pdf":
                # Handle PDF
                with st.spinner("Processing PDF..."):
                    images = convert_pdf_to_images(uploaded_file)
                    
                    if st.button("Extract Information"):
                        with st.spinner("Analyzing document..."):
                            all_results = []
                            for i, image in enumerate(images, 1):
                                result = analyze_image(image)
                                all_results.append(f"Page {i} Information:\n{result}")
                            
                            combined_results = "\n\n".join(all_results)
                            st.markdown(format_response(combined_results), unsafe_allow_html=True)
                            
            else:
                # Handle single image
                image = Image.open(uploaded_file)
                
                if st.button("Extract Information"):
                    with st.spinner("Analyzing document..."):
                        result = analyze_image(image)
                        st.markdown(format_response(result), unsafe_allow_html=True)

if __name__ == "__main__":
    main()