import altair as alt
import numpy as np
import pandas as pd
import streamlit as st
from dotenv import load_dotenv
import os
from mistralai import Mistral

from util import display_pdf, upload_pdf
from bon_livraison import extract_from_bl
from analyse_fourrage import extract_from_af
from document_classifier import classify_document

load_dotenv()

MISTRAL_API_KEY = os.environ.get("MISTRAL_API_KEY")

def main():
  """
  Main function to run the Streamlit app.
  """
  
  # Sidebar: Authentication for Mistral API
  if not MISTRAL_API_KEY:
    api_key = st.sidebar.text_input("Mistral API Key", type="password")
  else:
    api_key = MISTRAL_API_KEY
  
  if not api_key:
    st.warning("Enter API key to continue")
    return
    
  # Initialize Mistral API client
  client = Mistral(api_key=api_key)

  uploaded_file = st.file_uploader("Choisissez un PDF", type=["pdf"])

  document_source = None

  if uploaded_file:
    content = uploaded_file.read()
    preview_content = uploaded_file

    # Display the uploaded PDF
    display_pdf(content)
    
    # Prepare document source for OCR processing
    document_source = {
        "type": "document_url",
        "document_url": upload_pdf(client, content, uploaded_file.name)
    }
    content_type = "pdf"

  if document_source and st.button("Générer les données au format JSON"):
    # Process the document when the user clicks the button
    with st.spinner("Extracting JSON content..."):
      try:
        doc_type = classify_document(client, document_source)
        st.write(f"Document type: {doc_type}")

        if doc_type == "livraison":
          response = extract_from_bl(client, document_source)
          st.json(response)
        elif doc_type == "analyse":
          response = extract_from_af(client, document_source)
          st.json(response)
        else:
          st.error("Le document n'est pas supporté.")

      except Exception as e:
        # Display an error message if processing fails
        st.error(f"Processing error: {str(e)}")
  
if __name__ == "__main__":
  main()