import streamlit as st
import json
import copy
from receipt_gen_agent import receipt_agent
from PIL import Image
from huggingface_hub import HfApi
from os import listdir
from os.path import isfile, join
from io import BytesIO
from dotenv import load_dotenv
import os
import base64

load_dotenv()
token=os.getenv('hf_token')
api = HfApi()
repo_id = "wolf1997/receipt_scanner"  

agent=receipt_agent()


def capture_or_upload_image():
    """
    Provides options to capture or upload an image using Streamlit.
    
    Returns:
    PIL.Image.Image or None: The captured/uploaded image
    """
    st.sidebar.header("📸 Image Capture / Upload")
    
    # Choose method
    image_source = st.sidebar.radio(
        "Choose Image Source", 
        ["Upload", "Camera Capture"],
        key="image_source"
    )
    
    # Image upload
    if image_source == "Upload":
        uploaded_file = st.sidebar.file_uploader(
            "Upload Receipt Image", 
            type=['png', 'jpg', 'jpeg', 'gif'],
            help="Upload an image of the receipt",
            key="uploaded_file"
        )
        
        if uploaded_file is not None:
            # Convert to PIL Image
            pil_image = Image.open(uploaded_file)
            # Store in session state
            st.session_state.receipt_image = pil_image
            return pil_image
    
    # Camera capture
    else:
        captured_image = st.camera_input(
            "Take a picture of the receipt",
            key="camera_input"
        )
        
        if captured_image is not None:
            # Convert to PIL Image
            pil_image = Image.open(captured_image)
            # Store in session state
            st.session_state.receipt_image = pil_image
            return pil_image
    
    # Return image from session state if available
    return st.session_state.get("receipt_image", None)

def parse_receipt_data(receipt_str):
    """Parse the receipt dictionary from a string or dict."""
    if isinstance(receipt_str, str):
        try:
            return json.loads(receipt_str.replace("'", '"'))
        except json.JSONDecodeError:
            return {}
    return receipt_str

def receipt_editor_app():
    # Initialize session state variables if they don't exist
    if "receipt_dict" not in st.session_state:
        # Default initial receipt data
        initial_receipt_str = {
            "loc_name": "someplace",
            "table":"some table",
            "address": "somewhere",
            "date": "someday",
            "time": "sometime",
            "items": [
                {
                    "name": "something",
                    "price": 3.0,
                    "quantity": 1
                }
            ],
            "subtotal": 51.9,
            "tax": 4.68,
            "total": 56.58
        }
        st.session_state.receipt_dict = parse_receipt_data(initial_receipt_str)
        st.session_state.edited_receipt = copy.deepcopy(st.session_state.receipt_dict)
    
    if "receipt_image" not in st.session_state:
        st.session_state.receipt_image = None
        
    if "show_receipt_image" not in st.session_state:
        st.session_state.show_receipt_image = True
    
    # Set page title and layout
    st.set_page_config(page_title="Receipt Editor", layout="wide")

    # Capture or upload image
    receipt_image = capture_or_upload_image()
    
    if receipt_image is not None:
        # Option to display image
        show_image = st.sidebar.checkbox(
            "Show Receipt Image", 
            value=st.session_state.show_receipt_image,
            key="show_image_checkbox"
        )
        st.session_state.show_receipt_image = show_image
        
        if show_image:
            # Create two columns for image and details
            col1, col2 = st.sidebar.columns([1, 1])
            
            with col1:
                # Resize image to fit sidebar
                resized_image = receipt_image.copy()
                resized_image.thumbnail((600, 600))
                st.image(resized_image, caption="Uploaded Receipt")
            
            with col2:
                # Basic image details
                st.write(f"Image Size: {receipt_image.size}")
                st.write(f"Image Format: {receipt_image.format if hasattr(receipt_image, 'format') else 'Unknown'}")
                st.write(f"Color Mode: {receipt_image.mode}")
    
    if st.sidebar.button("Process Image"):
        initial_receipt_str = agent.receipt_gen(receipt_image)
        # Parse initial receipt data
        st.session_state.receipt_dict = parse_receipt_data(initial_receipt_str)
    
            
        # Update edited receipt in session state
        st.session_state.edited_receipt = copy.deepcopy(st.session_state.receipt_dict)
        st.sidebar.success("Receipt processed!")
        
        
    # Create a reference to session state for cleaner code
    edited_receipt = st.session_state.edited_receipt
    
    # Title
    st.title("Receipt Editor")
    
    # Basic Information Section
    st.header("Basic Information")
    col1, col2 = st.columns(2)
    with col1:
        edited_receipt['loc_name'] = st.text_input(
            "Location Name", 
            value=edited_receipt['loc_name'],
            key="loc_name"
        )
        edited_receipt['date'] = st.text_input(
            "Date", 
            value=edited_receipt['date'],
            key="date"
        )
        edited_receipt['table'] = st.text_input(
            "Table", 
            value=edited_receipt['table'],
            key="table"
        )
    
    with col2:
        edited_receipt['address'] = st.text_area(
            "Address", 
            value=edited_receipt['address'],
            key="address"
        )
        edited_receipt['time'] = st.text_input(
            "Time", 
            value=edited_receipt['time'],
            key="time"
        )
    
    # Items Section
    st.header("Items")
    
    # Expandable section for items
    with st.expander("Edit Items"):
        # Allow adding new items
        add_item = st.checkbox("Add New Item", key="add_item_checkbox")
        if add_item:
            with st.form(key="add_item_form"):
                new_item_name = st.text_input("New Item Name", key="new_item_name")
                new_item_price = st.number_input("New Item Price", min_value=0.0, step=0.01, key="new_item_price")
                new_item_quantity = st.number_input("New Item Quantity", min_value=1, step=1, value=1, key="new_item_quantity")
                
                submit_button = st.form_submit_button(label="Add to Receipt")
                if submit_button:
                    new_item = {
                        'name': new_item_name,
                        'price': new_item_price,
                        'quantity': new_item_quantity
                    }
                    edited_receipt['items'].append(new_item)
                    st.session_state.edited_receipt = edited_receipt
                    
        
    # Edit existing items
    for i, item in enumerate(edited_receipt['items']):
        with st.expander(f"Item {i+1}: {item['name']}"):
            col1, col2, col3 = st.columns(3)
            with col1:
                item['name'] = st.text_input(
                    f"Item {i+1} Name", 
                    value=item['name'], 
                    key=f"name_{i}"
                )
            with col2:
                item['price'] = st.number_input(
                    f"Item {i+1} Price", 
                    min_value=0.0, 
                    value=item['price'], 
                    step=0.01, 
                    key=f"price_{i}"
                )
            with col3:
                item['quantity'] = st.number_input(
                    f"Item {i+1} Quantity", 
                    min_value=1, 
                    value=item['quantity'], 
                    step=1, 
                    key=f"quantity_{i}"
                )
            
            # Option to remove item
            if st.button(f"Remove Item {i+1}", key=f"remove_{i}"):
                del edited_receipt['items'][i]
                st.session_state.edited_receipt = edited_receipt
               

    # Totals Section
    st.header("Totals")
    col1, col2, col3 = st.columns(3)
    edited_receipt=st.session_state.edited_receipt
    with col1:
        edited_receipt['subtotal'] = st.number_input(
            "Subtotal", 
            min_value=0.0, 
            value=edited_receipt['subtotal'], 
            step=0.01,
            key="subtotal"
        )
    
    with col2:
        edited_receipt['tax'] = st.number_input(
            "Tax", 
            min_value=0.0, 
            value=edited_receipt['tax'], 
            step=0.01,
            key="tax"
        )
    
    with col3:
        edited_receipt['total'] = st.number_input(
            "Total", 
            min_value=0.0, 
            value=edited_receipt['total'], 
            step=0.01,
            key="total"
        )

    # Recalculate totals button
    if st.button("Recalculate Totals", key="recalc_button"):
        edited_receipt=st.session_state.edited_receipt
        # Automatically calculate subtotal
        edited_receipt['subtotal'] = sum(item['price'] * item['quantity'] for item in edited_receipt['items'])
        
        # Recalculate total
        edited_receipt['total'] = edited_receipt['subtotal'] + edited_receipt['tax']
        
        # Update session state
        st.session_state.edited_receipt = edited_receipt
    
    # Display final edited receipt
    st.header("Edited Receipt")
    st.json(st.session_state.edited_receipt)

    # Handle file saving

    data_list = [f for f in listdir('new_receipt_data') if isfile(join('new_receipt_data', f))]
    if not data_list:
        st.session_state.data_list = []
    else:
        with open(f'new_receipt_data/{data_list[0]}', 'r') as openfile:
            # Reading from json file
            st.session_state.data_list = json.load(openfile)


    st.session_state.image_name = f'{len(st.session_state.data_list)+1}_new_receipt.jpg'
    
    # Option to save or export
    if st.button("Save Receipt Data", key="save_button"):
        try:
            # Create directories if they don't exist
            
            # Save receipt data
            data_list=st.session_state.data_list
            data_list.append({
                'receipt_name':st.session_state.image_name,
                'receipt_data': st.session_state.edited_receipt
            })
            
            # with open("new_receipt_data/receipt_data_list.json", "w") as f:
            #     json.dump(data_list, f, indent=2)
            json_string = json.dumps(data_list)
            bytes_data = json_string.encode('utf-8')
            api.upload_file(path_or_fileobj=bytes_data, repo_id=repo_id, repo_type="space",  path_in_repo="new_receipt_data/receipt_data_list.json",token=token)
            # Save image if available

            
            if st.session_state.receipt_image:
                # save_path = f'new_images/{st.session_state.image_name}'
                # st.session_state.receipt_image.save(save_path)
                buffered=BytesIO()
                image=st.session_state.receipt_image
                image.save(buffered, format='JPEG')
                # image_data = base64.b64encode(buffered.getvalue()).decode("utf-8")
                api.upload_file(path_or_fileobj=buffered.getvalue(), repo_id=repo_id,repo_type="space",  path_in_repo=f'new_images/{st.session_state.image_name}',token=token)
                st.success(f"Receipt data and image saved successfully!")
            else:
                st.success("Receipt data saved (no image available)")
                
        except Exception as e:
            st.error(f"Error saving receipt: {str(e)}")

    # Reset button to clear session state
    if st.sidebar.button("Reset All", key="reset_button"):
        # Clear specific session state values
        for key in list(st.session_state.keys()):
            del st.session_state[key]
        

def main():
    receipt_editor_app()

if __name__ == "__main__":
    main()