import streamlit as st
from PIL import Image
import os
import shutil
from ultralytics import YOLO
import tempfile
import pandas as pd
import plotly.express as px
from fetch_original import FileProcessor
from Damage_calculation import DamageCalculator
from generator import CSVGenerator
from pdf_report import PDFReportGenerator

class SegmentationModel:
    def __init__(self, model_path):
        self.model = YOLO(model_path)

    def predict(self, input_path):
        result = self.model.predict(source=input_path, save=True, show_conf=False, conf=0.70, save_txt=True)
        img_path = result[0].save_dir
        return img_path

class SegmentationApp:
    def __init__(self, car_parts_model_path, damage_model_path, output_folder):
        self.car_parts_model = SegmentationModel(car_parts_model_path)
        self.damage_model = SegmentationModel(damage_model_path)
        self.output_folder = output_folder

    def copy_folder(self, src_folder, dest_folder):
        if not os.path.exists(dest_folder):
            os.makedirs(dest_folder)
        for item in os.listdir(src_folder):
            s = os.path.join(src_folder, item)
            d = os.path.join(dest_folder, item)
            if os.path.isdir(s):
                shutil.copytree(s, d, dirs_exist_ok=True)
            else:
                shutil.copy2(s, d)

    def clean_output_folder(self):
        parts_output_folder = os.path.join(self.output_folder, 'parts')
        damage_output_folder = os.path.join(self.output_folder, 'damage')
        if os.path.exists(parts_output_folder):
            shutil.rmtree(parts_output_folder)
        if os.path.exists(damage_output_folder):
            shutil.rmtree(damage_output_folder)

    def run(self):
        st.title("Car Damage Analyses And Cost Estimation")
        st.markdown("### Important Guidelines for Using the App")
        st.write("""
        <div style="background-color: pink; color: black; padding: 10px; border-radius: 5px;">
        <ul>
            <li>Provide 4 images of the car from the following views: front, back, left, and right.</li>
            <li>Ensure the images are taken at proper angles, standing parallel to the car, and that they are clear.</li>
            <li>Make sure there are no obstacles between the car and the camera.</li>
            <li>The car should not be behind any objects in the images.</li>
        </ul>
        </div>
        """, unsafe_allow_html=True)

        uploaded_files = st.file_uploader("Choose up to 4 images...", type=["jpg", "jpeg", "png", "webp"], accept_multiple_files=True)

        if uploaded_files and len(uploaded_files) == 4:
            st.write("Running segmentation models on uploaded images...")
            car_parts_results_paths = []
            damage_results_paths = []

            self.clean_output_folder()

            with st.spinner("Analyzing total damage..."):
                for i, uploaded_file in enumerate(uploaded_files):
                    with tempfile.NamedTemporaryFile(delete=False, suffix='.jpg') as temp_file:
                        temp_file.write(uploaded_file.getvalue())
                        temp_image_path = temp_file.name

                    car_parts_results = self.car_parts_model.predict(temp_image_path)
                    car_parts_results_paths.append(car_parts_results)

                    damage_results = self.damage_model.predict(temp_image_path)
                    damage_results_paths.append(damage_results)

                self.car_parts_results_paths = car_parts_results_paths
                self.damage_results_paths = damage_results_paths

                for car_parts_results in car_parts_results_paths:
                    parts_output_folder = os.path.join(self.output_folder, 'parts')
                    self.copy_folder(car_parts_results, parts_output_folder)

                for damage_results in damage_results_paths:
                    damage_output_folder = os.path.join(self.output_folder, 'damage')
                    self.copy_folder(damage_results, damage_output_folder)

                file_processor = FileProcessor()
                file_processor.process_output_folder(self.output_folder)

                csv_generator = CSVGenerator(self.output_folder)
                damage_data = csv_generator.process_files()

                df = pd.read_csv(csv_generator.output_csv)
                st.write("Damage Estimation Report")
                st.dataframe(df)

                # Generate PDF report
                pdf_report = PDFReportGenerator(csv_generator.output_csv, self.output_folder, 'damage_estimation_report.pdf')
                pdf_report.generate_report()

                # # Add download button for the CSV file
                # csv_data = df.to_csv(index=False).encode('utf-8')
                # st.download_button(
                #     label="Download CSV",
                #     data=csv_data,
                #     file_name='damage_estimation.csv',
                #     mime='text/csv',
                # )

                # Add download button for the PDF report
                with open(pdf_report.pdf_path, 'rb') as f:
                    pdf_data = f.read()
                st.download_button(
                    label="Download PDF Report",
                    data=pdf_data,
                    file_name='damage_estimation_report.pdf',
                    mime='application/pdf',
                )

                damage_df = pd.DataFrame(damage_data)
                fig = px.pie(damage_df, values='coverage_percentage', names='damage_class', title='Total Damage Percentages by Type')
                st.plotly_chart(fig)

            st.success("Analysis Completed.")

        elif uploaded_files:
            st.error("Please upload exactly four images.")

if __name__ == "__main__":
    CAR_PARTS_MODEL_PATH = 'model/car_parts.pt'
    DAMAGE_MODEL_PATH = 'model/damage_4_class.pt'
    OUTPUT_FOLDER = 'Output'
    app = SegmentationApp(CAR_PARTS_MODEL_PATH, DAMAGE_MODEL_PATH, OUTPUT_FOLDER)
    app.run()