import os
import tempfile
import time

import cv2
import moviepy.editor as mp
import numpy as np
import streamlit as st
from streamlit_lottie import st_lottie
from tqdm import tqdm

from .models.deep_colorization.colorizers import eccv16
from .utils import load_lottieurl, format_time, colorize_frame, change_model

st.set_page_config(page_title="Image & Video Colorizer", page_icon="🎨", layout="wide")

loaded_model = eccv16(pretrained=True).eval()
current_model = "None"

col1, col2 = st.columns([1, 3])
with col1:
    lottie = load_lottieurl("https://assets5.lottiefiles.com/packages/lf20_RHdEuzVfEL.json")
    st_lottie(lottie)

with col2:
    st.write("""
    ## B&W Videos Colorizer
    ##### Upload a black and white video and get a colorized version of it.
    ###### ➠ This space is using CPU Basic so it might take a while to colorize a video.
    ###### ➠ If you want more models and GPU available please support this space by donating.""")


def main():
    model = st.selectbox(
        "Select Model (Both models have their pros and cons, I recommend trying both and keeping the best for your task)",
        ["ECCV16", "SIGGRAPH17"], index=0)

    loaded_model = change_model(current_model, model)
    st.write(f"Model is now {model}")

    uploaded_file = st.file_uploader("Upload your video here...", type=['mp4', 'mov', 'avi', 'mkv'])

    if st.button("Colorize"):
        if uploaded_file is None:
            st.warning("Please upload a video file first.")
        else:
            # Proceed with colorization
            file_extension = os.path.splitext(uploaded_file.name)[1].lower()
            if file_extension not in ['.mp4', '.avi', '.mov', '.mkv']:
                st.error("Unsupported file type. Please upload an MP4, AVI, MOV, or MKV file.")
                return # Stop execution if file type is wrong

            temp_video_path = None
            output_video_path = "output.mp4" # Consider making this temporary too
            converted_video_path = "converted_output.mp4" # Consider making this temporary too

            try:
                with tempfile.NamedTemporaryFile(delete=False, suffix=file_extension) as temp_file_obj:
                    temp_file_obj.write(uploaded_file.read())
                    temp_video_path = temp_file_obj.name

                st.write(f"Processing `{uploaded_file.name}`...")

                # Open the video using cv2.VideoCapture
                video = cv2.VideoCapture(temp_video_path)
                if not video.isOpened():
                    st.error("Could not open video file. It might be corrupted or an unsupported format.")
                    return

                fps = video.get(cv2.CAP_PROP_FPS)
                total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT))

                if fps == 0 or total_frames == 0:
                    st.error("Could not read video properties (FPS or frame count). The file might be invalid.")
                    video.release()
                    return

                # Attempt to extract audio
                try:
                    audio_clip = mp.AudioFileClip(temp_video_path)
                except Exception as e:
                    st.warning(f"Could not extract audio: {e}. Proceeding without audio.")
                    audio_clip = None

                col1, col2 = st.columns([0.5, 0.5])
                with col1:
                    st.markdown('<p style="text-align: center;">Before</p>', unsafe_allow_html=True)
                    st.video(temp_video_path)

                with col2:
                    st.markdown('<p style="text-align: center;">After</p>', unsafe_allow_html=True)

                    # Placeholder for the video display
                    video_display_placeholder = st.empty()
                    video_display_placeholder.info("Colorization in progress... Video will appear here.")

                    # Main processing
                    output_frames = []
                    progress_bar = st.progress(0)
                    time_text = st.text("Time Remaining: ")
                    start_time = time.time()

                    for i in tqdm(range(total_frames), unit='frame', desc="Colorizing Progress"):
                        ret, frame = video.read()
                        if not ret:
                            st.warning("Could not read all frames from the video.")
                            break

                        # colorize_frame now returns RGB, convert to BGR for VideoWriter
                        colorized_rgb_frame = colorize_frame(frame, loaded_model)
                        output_frames.append((colorized_rgb_frame * 255).astype(np.uint8))

                        elapsed_time = time.time() - start_time
                        frames_completed = len(output_frames)
                        if frames_completed > 0:
                            frames_remaining = total_frames - frames_completed
                            time_per_frame = elapsed_time / frames_completed
                            time_remaining = frames_remaining * time_per_frame
                            progress_bar.progress(frames_completed / total_frames)
                            if frames_completed < total_frames:
                                time_text.text(f"Time Remaining: {format_time(time_remaining)}")

                    if not output_frames:
                        st.error("No frames were colorized. Video processing failed.")
                        video.release()
                        return

                    time_text.empty()
                    progress_bar.empty()

                    with st.spinner("Merging frames and audio..."):
                        # Ensure frame_size is from the colorized frames (RGB)
                        # OpenCV VideoWriter expects (width, height)
                        frame_height, frame_width = output_frames[0].shape[:2]

                        # Output raw colorized video (without re-encoding audio yet)
                        # This will be RGB frames written as BGR
                        out = cv2.VideoWriter(output_video_path, cv2.VideoWriter_fourcc(*"mp4v"), fps, (frame_width, frame_height))
                        for frame_rgb in output_frames:
                            frame_bgr = cv2.cvtColor(frame_rgb, cv2.COLOR_RGB2BGR)
                            out.write(frame_bgr)
                        out.release()

                        # Now combine with audio and ensure Streamlit compatibility
                        final_clip = mp.VideoFileClip(output_video_path)
                        if audio_clip:
                            final_clip = final_clip.set_audio(audio_clip)

                        final_clip.write_videofile(converted_video_path, codec="libx264", audio_codec="aac")

                        video_display_placeholder.video(converted_video_path)
                        st.balloons()

                        with open(converted_video_path, "rb") as f_download:
                            st.download_button(
                                label="Download Colorized Video",
                                data=f_download.read(),
                                file_name=f"colorized_{uploaded_file.name}",
                                mime="video/mp4"
                            )

            except Exception as e:
                st.error(f"An error occurred during video processing: {e}")
            finally:
                if 'video' in locals() and video.isOpened():
                    video.release()
                # Clean up temporary files
                if temp_video_path and os.path.exists(temp_video_path):
                    os.unlink(temp_video_path)
                if os.path.exists(output_video_path): # output_video_path = "output.mp4"
                    os.unlink(output_video_path)
                # converted_video_path is what's displayed and downloaded, so it might be kept
                # or also made temporary and read into memory for download.
                # For simplicity here, we're not deleting converted_video_path immediately
                # but in a real scenario, it should also be managed (e.g. via tempfile)

if __name__ == "__main__":
    main()
    st.markdown(
        "###### Made with :heart: by [Clément Delteil](https://www.linkedin.com/in/clementdelteil/) [![this is an "
        "image link](https://i.imgur.com/thJhzOO.png)](https://www.buymeacoffee.com/clementdelteil)")
    st.markdown(
        "###### [Blog post of the project](https://medium.com/geekculture/creating-a-web-app-to-colorize-images-and-youtube-videos-80f5be2d0f68)"
    )