import streamlit as st
import os
import glob
import base64
import json
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.image as mpimg
from langchain_openai import ChatOpenAI
from langchain_core.pydantic_v1 import BaseModel, Field
from langchain_core.messages import HumanMessage, SystemMessage
from langchain_core.output_parsers import JsonOutputParser
from langchain_core.runnables import chain
from PIL import Image as PILImage
from io import BytesIO

# Streamlit title
st.title("🚗 Vehicle Information Extraction from Images 🚙")

# Custom CSS Styling for the app
st.markdown("""
    <style>
        body {
            background-color: #f0f8ff;
            font-family: 'Arial', sans-serif;
            color: #333;
        }

        .stButton button {
            background-color: #1E90FF;
            color: white;
            border-radius: 5px;
            font-size: 16px;
            padding: 10px 20px;
            transition: transform 0.2s ease-in-out;
        }

        .stButton button:hover {
            background-color: #4682b4;
            transform: scale(1.05);
        }

        .stTitle {
            font-size: 36px;
            font-weight: bold;
            color: #1E90FF;
        }

        .stSubheader {
            color: #4682b4;
            font-size: 24px;
            font-weight: 600;
            text-align: center;
        }

        .stImage {
            border-radius: 10px;
            box-shadow: 0 4px 8px rgba(0, 0, 0, 0.1);
        }

        .stJson {
            background-color: #f8f9fa;
            padding: 15px;
            border-radius: 5px;
            box-shadow: 0 4px 8px rgba(0, 0, 0, 0.1);
            font-size: 16px;
        }

        .grid-container {
            display: grid;
            grid-template-columns: repeat(3, 1fr);
            gap: 10px;
            padding: 10px;
        }

        .grid-container img {
            border-radius: 10px;
            box-shadow: 0 4px 8px rgba(0, 0, 0, 0.1);
            transition: transform 0.2s ease-in-out;
        }

        .grid-container img:hover {
            transform: scale(1.05);
        }

        @keyframes fadeIn {
            from {
                opacity: 0;
            }
            to {
                opacity: 1;
            }
        }

        .fade-in {
            animation: fadeIn 1s ease-in-out;
        }

    </style>
""", unsafe_allow_html=True)

# Image display function with animation
def display_image_grid(image_paths, rows=2, cols=3, figsize=(10, 7)):
    fig = plt.figure(figsize=figsize)
    max_images = rows * cols
    image_paths = image_paths[:max_images]

    for idx, path in enumerate(image_paths):
        ax = fig.add_subplot(rows, cols, idx + 1)
        img = mpimg.imread(path)
        ax.imshow(img)
        ax.axis('off')
        filename = path.split('/')[-1]
        ax.set_title(filename)

    plt.tight_layout()
    st.pyplot(fig)

# Image encoding function
def image_encoding(inputs):
    """Load and convert image to base64 encoding"""
    with open(inputs["image_path"], "rb") as image_file:
        image_base64 = base64.b64encode(image_file.read()).decode("utf-8")
    return {"image": image_base64}

# Streamlit Interface for uploading images and showing results
st.header("🚗 Upload Vehicle Images for Information Extraction 🚙")

# Option to select either single or batch image upload
upload_option = st.radio("Select Upload Type", ["Single Image Upload", "Batch Images Upload"])

# Single Image Upload
if upload_option == "Single Image Upload":
    st.subheader("Upload a Single Vehicle Image")
    uploaded_image = st.file_uploader("Choose an Image (JPEG, PNG, GIF, BMP, etc.)", type=["jpeg", "png", "gif", "bmp", "jpg"])

    if uploaded_image is not None:
        # Display the uploaded image with animation
        image = PILImage.open(uploaded_image)
        st.image(image, caption="Uploaded Image", use_container_width=True, class_="fade-in")

        # Convert the uploaded image to base64
        image_path = "/tmp/uploaded_image" + os.path.splitext(uploaded_image.name)[1]
        with open(image_path, "wb") as f:
            f.write(uploaded_image.getbuffer())

        # Add button to trigger information extraction with animated hover effect
        if st.button("Extract Vehicle Information"):
            # Process the image through the pipeline
            output = pipeline.invoke({"image_path": image_path})

            # Show the results in a user-friendly format
            st.subheader("Extracted Vehicle Information", class_="fade-in")
            st.json(output, class_="fade-in")

# Batch Images Upload
elif upload_option == "Batch Images Upload":
    st.sidebar.header("Batch Image Upload")
    batch_images = st.sidebar.file_uploader("Upload Images (JPEG, PNG, GIF, BMP, etc.)", type=["jpeg", "png", "gif", "bmp", "jpg"], accept_multiple_files=True)

    if batch_images:
        batch_input = [{"image_path": f"/tmp/{file.name}"} for file in batch_images]
        for file in batch_images:
            with open(f"/tmp/{file.name}", "wb") as f:
                f.write(file.getbuffer())

        # Add button to trigger batch information extraction
        if st.button("Extract Vehicle Information from Batch"):
            # Process the batch and display the results in a DataFrame
            batch_output = pipeline.batch(batch_input)
            df = pd.DataFrame(batch_output)
            st.dataframe(df, use_container_width=True)

        # Display batch images in a grid with animation
        st.subheader("Images in Grid", class_="fade-in")
        image_paths = [f"/tmp/{file.name}" for file in batch_images]
        st.markdown('<div class="grid-container">', unsafe_allow_html=True)
        for image_path in image_paths:
            st.markdown(f'<img src="data:image/jpeg;base64,{base64.b64encode(open(image_path, "rb").read()).decode()}" class="fade-in"/>', unsafe_allow_html=True)
        st.markdown('</div>', unsafe_allow_html=True)