import streamlit as st
from utils.logger import setup_logger
import pandas as pd
from PIL import Image
import os

logger = setup_logger(__name__)


def prune_bank_statement_for_display(analysis_results_for_id):
    data_to_display = {}

    data_to_display["document_category"] = "bank_statement"
    data_to_display["document_type"] = "bank_statement"

    data_to_display["account_holder_name"] = analysis_results_for_id.get(
        "account_holder_name", None)
    data_to_display["account_holder_address"] = analysis_results_for_id.get(
        "account_holder_address", None)
    data_to_display["bank_name"] = analysis_results_for_id.get(
        "bank_name", None)
    data_to_display["account_number"] = analysis_results_for_id.get(
        "account_number", None)
    data_to_display["sort_code"] = analysis_results_for_id.get(
        "sort_code", None)
    data_to_display["statement_start_date"] = analysis_results_for_id.get(
        "statement_start_date", None)
    data_to_display["statement_end_date"] = analysis_results_for_id.get(
        "statement_end_date", None)
    data_to_display["salary_credits"] = analysis_results_for_id.get(
        "salary_credits", None)

    return data_to_display


def display_bank_statement(extracted_files, analysis_results_pruned):

    col1, col2 = st.columns([2, 3])

    logger.info(f"file_path while displaying: {extracted_files}")
    st.markdown("---")

    with col1:
        if len(extracted_files) > 1:
            st.image(extracted_files, caption=[os.path.basename(
                img) for img in extracted_files], use_container_width=True)
        else:
            image = Image.open(extracted_files[0])
            st.image(image, caption=os.path.basename(
                extracted_files[0]))  # ,
            # use_container_width=True)

        logger.info(
            f"analysis_results_pruned : {analysis_results_pruned}")

    with col2:

        dict_str = {}

        for key, value in analysis_results_pruned.items():
            if key != 'salary_credits':
                dict_str[key] = value

        simple_df = pd.DataFrame.from_dict(
            dict_str,
            orient='index', columns=['Value']).reset_index()
        simple_df.columns = ['Key', 'Value']
        simple_df = simple_df.fillna(value="Missing")
        simple_df.index += 1
        st.dataframe(simple_df, use_container_width=True)

        st.markdown("Salary Credits")

        salary_dict = analysis_results_pruned['salary_credits']
        logger.info(f"salary_dict : {salary_dict}")
        for salary_details in salary_dict:
            simple_df = pd.DataFrame.from_dict(
                salary_details,
                orient='index', columns=['Value']).reset_index()
            simple_df.columns = ['Key', 'Value']
            simple_df = simple_df.fillna(value="Missing")
            simple_df.index += 1
            st.dataframe(simple_df, use_container_width=True)

        logger.info(f"simple_df: {simple_df}")