import streamlit as st
import uuid
from datetime import datetime
from utils.session_state import reset_state
from utils.process_files import process_uploaded_files
from utils.document_display import display_based_on_card
from llm.document_analyzer import analyze_files
import json
from utils.logger import setup_logger

logger = setup_logger(__name__)


def upload_documents():
    uploaded_files = st.file_uploader(
        "Upload Images, PDFs, or ZIP files",
        type=["png", "jpg", "jpeg", "pdf", "zip"],
        accept_multiple_files=True,
        key="file_uploader"
    )

    if uploaded_files and st.session_state.get('current_upload') is None:
        # if uploaded_files:
        new_upload_id = str(uuid.uuid4())[:8]  # Shorter UUID for readability
        new_upload_id = "_".join([str(datetime.now()), new_upload_id])
        reset_state(upload_id=new_upload_id)
        st.session_state['analyze_clicked'] = False

        # Process files
        file_paths, file_groups, temp_dir, application_form, memo = process_uploaded_files(
            uploaded_files)
        st.session_state['uploads'][new_upload_id]['file_groups'] = file_groups
        st.session_state['uploads'][new_upload_id]['temp_dir'] = temp_dir
        st.session_state['uploads'][new_upload_id]['application_form'] = application_form
        st.session_state['uploads'][new_upload_id]['memo'] = memo

    # Get current upload
    current_upload = st.session_state.get('current_upload')

    if current_upload:
        current_data = st.session_state['uploads'][current_upload]
        file_groups = current_data['file_groups']
        temp_dir = current_data['temp_dir']

        if file_groups:

            analyze_clicked = st.button("🔍 Analyze")

            if analyze_clicked:
                st.session_state['analyze_clicked'] = True
                reset_state()  # Reset raw & display results only

        # Analysis step: only when user has clicked Analyze
        if st.session_state.get('analyze_clicked', False):

            if current_data['values_raw'] is None:
                with st.spinner("Analyzing documents..."):
                    analysis_results_groups, json_output_path = analyze_files(
                        file_groups=file_groups,
                        temp_dir=temp_dir,
                        current_upload=current_upload
                    )
                    current_data['values_raw'] = analysis_results_groups

            if current_data['values_display'] is None:
                current_data['values_display'] = {}

            for original_file, extracted_files in file_groups.items():
                display_based_on_card(
                    original_file=original_file,
                    analysis_results_for_original_file=current_data['values_raw'][original_file],
                    extracted_files=extracted_files,
                    current_upload=current_upload
                )

            st.download_button(
                label="Download Analysis JSON",
                data=json.dumps(
                        # current_data['values_raw'],
                        current_data['results_transformed'],
                        indent=4),
                file_name="analysis_results.json",
                mime="application/json"
            )