import streamlit as st
from PIL import Image
import os
from uuid import uuid4
import fitz
from invoice_generator import generate_invoice

st.set_page_config(page_title="Invoice generator", layout="wide")
output_folder = "output"
template = "template.tex"


def get_image_from_pdf(pdf_path):
    doc = fitz.open(pdf_path)
    page = doc[0]
    mat = fitz.Matrix(2, 2)
    pix = page.get_pixmap(matrix=mat)
    img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
    return img

def display_invoice(image_path):
    output_pdf = "invoice_" + os.path.basename(image_path).split(".")[0] + ".pdf"
    path_to_output_pdf = f"{output_folder}/{output_pdf}"
    try:
        status = generate_invoice(image_path, output_pdf, template, output_folder)
        if status == 0:
            st.write("Image is irrelevant, upload another one")
            return
        print(f"Generated invoice: {path_to_output_pdf}")
        st.query_params["generated_pdf"] = path_to_output_pdf
        return get_image_from_pdf(path_to_output_pdf)
    except Exception as e:
        st.write("Could not generate invoice, please try again")
        print(e)
        return None


st.title("Upload FNOL photo")

col1, col2, col3 = st.columns([4, 1, 4])

with col1:
    uploaded_image = st.file_uploader("Upload photo", type=["jpg", "jpeg", "png"])
    if uploaded_image:
        try:
            image = Image.open(uploaded_image)
            image_path = f"{output_folder}/{str(uuid4())[:5]}.png"
            image.save(image_path)
            print(f"Image: {image_path}")
            st.image(image, caption="Uploaded photo", width=300)
            st.query_params["image"] = image_path
        except Exception as e:
            st.write(f"Coudn't load image: {e}")

with col2:
    if st.query_params.get("image"):
        if st.button("Generate invoice"):
            with st.spinner("Generating..."):
                st.session_state["invoice"] = display_invoice(st.query_params["image"])
                if st.session_state["invoice"]:
                    st.query_params["status"] = "loaded"
    else:
        st.button("Generate invoice", disabled=True)
    if st.query_params.get("generated_pdf"):
        with open(st.query_params["generated_pdf"], "rb") as f:
            file_data = f.read()
        st.download_button(
            label="Download invoice",
            data=file_data,
            file_name="generated_invoice.pdf",
            mime="application/pdf"
        )

with col3:
    if st.query_params.get("status") == "loaded":
        st.image(st.session_state["invoice"], caption="Generated invoice", use_column_width=True)