from itertools import product
import streamlit as st
import numpy as np
import pandas as pd
from PIL import Image, ImageOps
import time
from paddleocr import PaddleOCR
import os
from dotenv import load_dotenv
from huggingface_hub import login

load_dotenv()  # Load .env file
huggingface_token = os.getenv("HF_TOKEN")
login(huggingface_token)

##########################LLAMA3BI################################
from huggingface_hub import InferenceClient
client = InferenceClient(api_key=huggingface_token)
messages = [
    {"role": "system", "content": """Your task is to get the product details out of the text given. 
                                    The text given will be raw text from OCR of social media images of products,
                                    and the goal is to get product details and description so that it can be used for e-commerce product listings. 
                                    TRY TO KEEP THE LISTING IN FOLLOWING FORMAT. 
                                    📦 [Product Name]
                                    💰 Price: $XX.XX
                                    ✨ Key Features:
                                    •⁠  ⁠[Main Feature 1]
                                    •⁠  ⁠[Main Feature 2]
                                    •⁠  ⁠[Main Feature 3]
                                    📸 [Product Image]
                                    🏷 Available Now
                                    ✈️ Prime Shipping Available
                                    🛍 Shop Now: [Link]
                                    🔍 Search: [Main Keywords]
                                    [#RelevantHashtags] """},
]

# Initialize PaddleOCR model
ocr = PaddleOCR(use_angle_cls=True, lang='en')

# Team details
team_members = [
    {"name": "Aman Deep", "image": "aman.jpg"},  # Replace with actual paths to images
    {"name": "Nandini", "image": "nandini.jpg"},
    {"name": "Abhay Sharma", "image": "abhay.jpg"},
    {"name": "Ratan Prakash Mishra", "image": "anandimg.jpg"}
]

# Function to preprocess images for the model
def preprocess_image(image):
    """
    Preprocess the input image for model prediction.
    Args:
        image (PIL.Image): Input image in PIL format.
    Returns:
        np.ndarray: Preprocessed image array ready for prediction.
    """
    try:
        # Resize image to match model input size
        img = image.resize((128, 128), Image.LANCZOS)

        # Convert image to NumPy array
        img_array = np.array(img)

        # Check if the image is grayscale and convert to RGB if needed
        if img_array.ndim == 2:  # Grayscale image
            img_array = np.stack([img_array] * 3, axis=-1)
        elif img_array.shape[2] == 1:  # Single-channel image
            img_array = np.concatenate([img_array, img_array, img_array], axis=-1)

        # Normalize pixel values to [0, 1] range
        img_array = img_array / 255.0

        # Add batch dimension
        img_array = np.expand_dims(img_array, axis=0)

        return img_array

    except Exception as e:
        print(f"Error processing image: {e}")
        return None

# Function to display team members in circular format
def display_team_members(members, max_members_per_row=4):
    num_members = len(members)
    num_rows = (num_members + max_members_per_row - 1) // max_members_per_row

    for i in range(num_rows):
        cols = st.columns(min(max_members_per_row, num_members - i * max_members_per_row))
        for j, member in enumerate(members[i * max_members_per_row:(i + 1) * max_members_per_row]):
            with cols[j]:
                img = Image.open(member["image"])
                st.image(img, use_column_width=True)
                st.write(member["name"])

# Function to simulate loading process with a progress bar
def simulate_progress():
    progress_bar = st.progress(0)
    for percent_complete in range(100):
        time.sleep(0.02)
        progress_bar.progress(percent_complete + 1)

# Title and description
st.title("Product Listing Assistant")

# Navbar with task tabs
st.sidebar.title("Navigation")
st.sidebar.write("Team Name: Sadhya")
app_mode = st.sidebar.selectbox("Choose the task", ["Welcome", "Project Details", "Team Details", "Extract Product Details"])

if app_mode == "Welcome":
    st.write("# Welcome to the Product Listing Assistant! 🎉")
elif app_mode == "Project Details":
    st.write("""
    ## Project Overview:
    - Automates product listings from social media content.
    - Extracts product details from posts using OCR and LLMs.
    - Outputs structured, engaging, and optimized e-commerce listings.
    """)
elif app_mode == "Team Details":
    st.write("## Meet Our Team:")
    display_team_members(team_members)
elif app_mode == "Extract Product Details":
    st.write("## Extract Product Details Using OCR and LLM")
    post_url = st.text_input("Enter Post URL:")
    uploaded_files = st.file_uploader("Upload Product Images", type=["jpeg", "png", "jpg"], accept_multiple_files=True)

    if post_url:
        st.write("### Processed Details:")
        # Add Instagram post processing logic here.

    if uploaded_files:
        st.write("### Uploaded Images:")
        simulate_progress()
        for uploaded_image in uploaded_files:
            image = Image.open(uploaded_image)
            st.image(image, use_column_width=True)
            simulate_progress()
            st.write("Details extracted:")
            # Add OCR and LLM processing logic here.