import streamlit as st
from groq import Groq
import base64
import os
import io
from PIL import Image
import traceback
from datetime import datetime
import shutil  # For moving files
import json  # For JSON formatting

# List of available models
MODELS = [
    "llama-3.2-90b-vision-preview",
    "llama-3.2-11b-vision-preview",
    "llava-v1.5-7b-4096-preview"
]

# Folder paths
RECEIPT_FOLDER = "/home/project_murky/Interactive-Image-Analysis-Tool/upload_receipts"
COMPLETED_FOLDER = "/home/project_murky/Interactive-Image-Analysis-Tool/completed_receipts"
PROMPT_FILE = "/home/project_murky/Interactive-Image-Analysis-Tool/receipt_prompt.txt"

def get_image_list():
    """Retrieve list of image file paths from the receipt folder."""
    return [f for f in os.listdir(RECEIPT_FOLDER) if f.endswith(('.png','JPG', 'jpg', '.jpeg'))]

def load_image(image_name):
    """Load an image from the specified folder."""
    image_path = os.path.join(RECEIPT_FOLDER, image_name)
    return Image.open(image_path), image_path

def compress_image(image, max_size=(800, 800), quality=95):
    img = Image.open(image) if isinstance(image, str) else image
    img.thumbnail(max_size)
    buffered = io.BytesIO()
    img.save(buffered, format="JPEG", quality=quality)
    return buffered.getvalue()

def encode_image(image):
    if isinstance(image, Image.Image):
        buffered = io.BytesIO()
        image.save(buffered, format="JPEG", quality=95)
        return base64.b64encode(buffered.getvalue()).decode('utf-8')
    else:
        compressed = compress_image(image)
        return base64.b64encode(compressed).decode('utf-8')

def create_client():
    api_key = "<TRUFFLEHOG_REDACTED_GROQ>"
    return Groq(api_key=api_key)

def read_prompt():
    """Read the prompt from the specified file."""
    with open(PROMPT_FILE, 'r') as file:
        return file.read().strip()

def analyze_input(text_input, Quick_Input):
    if Quick_Input == "Input Manually":
        return text_input.strip()
    elif Quick_Input == "Analyze Receipt":
        return "Project Overview: Not every receipt has every column, use what data you have do not make anything up this is serious shit and is related to my business. I need exact details or do not provide anything and say you do not know. So lets say the receipt doesnt show the tax and just has a total. Then we would have to go backwards and we would have to take the total and divide it by .13 and go from there please. Receipt Extraction Objective: Develop a system to extract key information from receipts in a structured format, enabling efficient data processing and analysis for applications like expense tracking and financial reporting. Tasks: Define Data Structure: Create a template outlining data points to be extracted from receipts, including: Vendor Name, Total Amount, Date of Purchase, Items Bought, Taxes Applied, Discounts Applied. Implement Extraction Logic: Develop algorithms or utilize existing libraries to extract the defined fields from receipt images or text. Validation: Ensure data accuracy by validating totals, checking for required fields, and managing errors. Testing: Conduct thorough testing with various receipt formats to ensure robustness and accuracy of the extraction process. Documentation: Document the code, the extraction process, and instructions for system use to facilitate future reference and maintenance. Goals: Create a reliable and efficient receipt extraction system that handles different formats and layouts. Maintain data integrity by ensuring all extracted values are calculated and validated correctly. Provide clear, user-friendly output for easy reporting and analysis."
    elif Quick_Input == "Text in Image":
        return "What does the text in this photo say? Respond in JSON format."
    elif Quick_Input == "Image Reasoning":
        return "Let's work this out in a step-by-step way to ensure we have the right answer. Respond in JSON format."
    else:
        return text_input.strip()

def save_response(response):
    directory = "/home/project_murky/Interactive-Image-Analysis-Tool/receipts"
    os.makedirs(directory, exist_ok=True)
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    filename = f"response_{timestamp}.txt"
    filepath = os.path.join(directory, filename)
    
    with open(filepath, 'w') as f:
        f.write(response)
    
    return filepath

def move_to_completed(image_path):
    """Move the image to the completed folder."""
    if os.path.exists(image_path):
        shutil.move(image_path, COMPLETED_FOLDER)

def process_image_and_text(image, text_input, Quick_Input, model, image_path=None):
    st.info("Image is being analyzed, please wait a moment...")
    if Quick_Input == "Input Manually" and not text_input.strip():
        return {"error": "Please enter a question or select a quick input option!"}, ""
    
    text_input = text_input.strip()
    client = create_client()
    base64_image = encode_image(image)
    prompt_text = read_prompt()  # Read prompt from file
    input_text = analyze_input(text_input, Quick_Input)
    
    # Prepare the full prompt for the model, clearly requesting a JSON response
    full_prompt = f"{prompt_text}\n\nInstructions: Please respond in JSON format, strictly adhering to the structure outlined in the prompt above.\n\nUser Input: {input_text}"
    
    try:
        chat_completion = client.chat.completions.create(
            messages=[{
                "role": "user",
                "content": [
                    {"type": "text", "text": full_prompt},
                    {
                        "type": "image_url",
                        "image_url": {
                            "url": f"data:image/jpeg;base64,{base64_image}",
                        },
                    },
                ],
            }],
            model=model,
            temperature=1,
        )
        st.success("Vision model has completed the response.")
        response_content = chat_completion.choices[0].message.content.strip()
        
        # Convert the response to a formatted JSON string
        try:
            response_json = json.loads(response_content)  # Parse JSON
            formatted_response = json.dumps(response_json, indent=4)  # Pretty print JSON
        except json.JSONDecodeError:
            formatted_response = json.dumps({"error": "Invalid JSON response received."})  # Fallback to a JSON error message
        
        filepath = save_response(formatted_response)
        
        # Move image to completed folder if processed from dropdown
        if image_path:
            move_to_completed(image_path)
        
        return formatted_response, f"Response saved to: {filepath}"
    except Exception as e:
        error_traceback = traceback.format_exc()
        print(f"Error occurred: {error_traceback}")
        return json.dumps({"error": f"Error occurred: {str(e)}"}), ""

def update_textbox_based_on_quick_input(Quick_Input):
    return analyze_input("", Quick_Input)

def move_image(image_name):
    """Handle the move image action."""    
    image_path = os.path.join(RECEIPT_FOLDER, image_name)
    move_to_completed(image_path)
    return f"Image {image_name} has been moved to completed folder."

# Streamlit UI setup
st.title("Image Analysis and Reasoning - Large Vision Model")
st.markdown("### Project by Scott")

# Dropdown for selecting images
image_list = get_image_list()
image_name = st.selectbox("Select Image from Receipts Folder", image_list)
image, image_path = load_image(image_name) if image_name else (None, None)

# Display the selected image
if image:
    st.image(image, caption=image_name, use_column_width=True)

# Dropdown for selecting models
model_select = st.selectbox("Select Vision Model", MODELS)

# Clear button functionality
if st.button("Clear Answer and Image"):
    st.session_state.response = ""
    st.session_state.image = None
    st.rerun()

# Quick Input dropdown
Quick_Input = st.selectbox(
    "Quick Input",
    ["Input Manually", "Analyze Receipt", "Text in Image", "Image Reasoning"],
    index=0
)

# Text input box
text_input = st.text_area("Enter Question", placeholder="Please enter your question...", height=150)

# Update text_input based on Quick Input
if Quick_Input:
    updated_text = update_textbox_based_on_quick_input(Quick_Input)
    text_input = st.text_area("Enter Question", value=updated_text, height=150)

# Submit button
if st.button("Submit"):
    if image and (Quick_Input != "Input Manually" or text_input.strip()):
        response_content, file_info = process_image_and_text(image, text_input, Quick_Input, model_select, image_path)
        
        # Display response in chat style
        st.markdown("### Vision Model Response:")
        st.markdown(f"<div style='border: 1px solid #ccc; padding: 10px; border-radius: 5px;'>{response_content}</div>", unsafe_allow_html=True)
        if file_info:
            st.markdown(file_info)

# Additional actions
if st.button("Move Image to Completed"):
    if image_name:
        move_message = move_image(image_name)
        st.success(move_message)


refactor this code please