import logging
import time

import numpy as np
import rembg
import torch
from PIL import Image
from rotate import rotate
import streamlit as st

import sys
import os


from tsr.system import TSR
from x3D_utils import remove_background, resize_foreground

import logging
import time
import streamlit as st
import torch
from datetime import datetime


# Hàm tùy chỉnh để hiển thị thông báo kèm thời gian tương tự logging
def st_info_with_logging_format(message):
    current_time = datetime.now().strftime('%Y-%m-%d %H:%M:%S:%f')[:-3]  # Định dạng thời gian giống logging
    formatted_message = f"{current_time} - INFO - {message}"
    st.info(formatted_message)

def st_success_with_logging_format(message):
    current_time = datetime.now().strftime('%Y-%m-%d %H:%M:%S:%f')[:-3]  # Định dạng thời gian giống logging
    formatted_message = f"{current_time} - INFO - {message}"
    st.success(formatted_message)

# Thay thế các thông báo st.info bằng hàm mới
def convert_to_3d(image_path, output_filename='', isHuman=False, isCloth=False, cloth_cat=''):
    with st.expander("ImageTo3D Extract Infomation"):
        class Timer:
            def __init__(self):
                self.items = {}
                self.time_scale = 1000.0  # ms
                self.time_unit = "seconds"

            def start(self, name: str) -> None:
                if torch.cuda.is_available():
                    torch.cuda.synchronize()
                self.items[name] = time.time()
                st_info_with_logging_format(f"{name} ...")

            def end(self, name: str) -> float:
                if name not in self.items:
                    return
                if torch.cuda.is_available():
                    torch.cuda.synchronize()
                start_time = self.items.pop(name)
                delta = time.time() - start_time
                t = delta * self.time_scale
                st_success_with_logging_format(f"{name} finished in {(t / 1000):.2f} {self.time_unit}.")

        timer = Timer()

        logging.basicConfig(
            format="%(asctime)s - %(levelname)s - %(message)s", level=logging.INFO
        )

        if not torch.cuda.is_available():
            device = "cpu"
        else:
            device = "cuda:0"

        timer.start("Initializing model")
        model = TSR.from_pretrained(
            config_path=r"config.yaml",
            weight_path=r"model.ckpt"
        )

        model.renderer.set_chunk_size(10_000)  # 0 for no chunking; default is 8192
        model.to(device)
        timer.end("Initializing model")

        timer.start("Removing background")
        if isHuman:
            rembg_session = rembg.new_session(model_name="u2net_human_seg")
        elif isCloth:
            rembg_session = rembg.new_session(model_name="u2net_cloth_seg")
        else:
            rembg_session = rembg.new_session()

        if isCloth and cloth_cat != '':
            image = remove_background(Image.open(image_path), rembg_session, cloth_category=cloth_cat)
        else:
            image = remove_background(Image.open(image_path), rembg_session)
        image = resize_foreground(image, 0.85)
        image = np.array(image).astype(np.float32) / 255.0
        image = image[:, :, :3] * image[:, :, 3:4] + (1 - image[:, :, 3:4]) * 0.5
        image = Image.fromarray((image * 255.0).astype(np.uint8))

        timer.end("Removing background")

        timer.start("Running model on image")
        with torch.no_grad():
            scene_codes = model([image], device=device)
        timer.end("Running model on image")

        timer.start("Extracting mesh")
        mesh = model.extract_mesh(scene_codes, resolution=256)[0]
        timer.end("Extracting mesh")

        timer.start("Rotating object")
        mesh = rotate(mesh)
        timer.end("Rotating object")

        timer.start("Saving generated object")
        if output_filename == '':
            output_filename = f"{image_path.split('.')[-2]}_out"
        output_filepath = f"{output_filename}.glb"
        mesh.export(output_filepath)
        timer.end("Saving generated object")

        return output_filepath