import os
import cv2
import gradio as gr
import AnimeGANv3_src
import numpy as np
import logging
import onnxruntime as ort
import tempfile
import threading

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

# Check ONNX Runtime providers
AVAILABLE_PROVIDERS = ort.get_available_providers()

logging.info(
    "Available ONNX Runtime providers: %s",
    AVAILABLE_PROVIDERS
)

if "CUDAExecutionProvider" in AVAILABLE_PROVIDERS:
    logging.info("CUDA provider is available.")
else:
    logging.warning(
        "CUDA provider is NOT available. "
        "The model may run on CPU."
    )


class AnimeGANv3:

    def __init__(self):
        os.makedirs("output", exist_ok=True)
        os.makedirs("frames", exist_ok=True)

        self.lock = threading.Lock()

        logging.info("AnimeGANv3 initialized.")

    def process_frame(self, frame, style_code, det_face):
        """
        Convert a single OpenCV BGR frame
        using the existing AnimeGANv3 module.
        """

        frame_rgb = cv2.cvtColor(
            frame,
            cv2.COLOR_BGR2RGB
        )

        output = AnimeGANv3_src.Convert(
            frame_rgb,
            style_code,
            det_face
        )

        if output is None:
            raise RuntimeError(
                "AnimeGANv3 returned an empty frame."
            )

        if not isinstance(output, np.ndarray):
            output = np.asarray(output)

        if output.ndim != 3 or output.shape[2] != 3:
            raise RuntimeError(
                f"Unexpected output shape: {output.shape}"
            )

        # RGB -> BGR
        return cv2.cvtColor(
            output,
            cv2.COLOR_RGB2BGR
        )

    def inference(self, video_path, style, if_face="No"):

        if not video_path:
            raise gr.Error("Please upload a video first.")

        style_codes = {
            "AnimeGANv3_Arcane": "A",
            "AnimeGANv3_Trump v1.0": "T",
            "AnimeGANv3_Shinkai": "S",
            "AnimeGANv3_PortraitSketch": "P",
            "AnimeGANv3_Hayao": "H",
            "AnimeGANv3_Disney v1.0": "D",
            "AnimeGANv3_JP_face v1.0": "J",
            "AnimeGANv3_Kpop v2.0": "K",
        }

        style_code = style_codes.get(style)

        if style_code is None:
            raise gr.Error("Invalid AnimeGANv3 style.")

        det_face = if_face == "Yes"

        logging.info(
            "Starting video conversion: style=%s, face_detection=%s",
            style,
            det_face
        )

        with self.lock:

            cap = None
            out = None
            save_path = None

            try:
                cap = cv2.VideoCapture(video_path)

                if not cap.isOpened():
                    raise RuntimeError(
                        "Could not open the uploaded video."
                    )

                fps = cap.get(cv2.CAP_PROP_FPS)

                if not np.isfinite(fps) or fps <= 0:
                    fps = 25.0

                frame_count = int(
                    cap.get(cv2.CAP_PROP_FRAME_COUNT)
                )

                width = int(
                    cap.get(cv2.CAP_PROP_FRAME_WIDTH)
                )

                height = int(
                    cap.get(cv2.CAP_PROP_FRAME_HEIGHT)
                )

                if width <= 0 or height <= 0:
                    raise RuntimeError(
                        "Could not determine video dimensions."
                    )

                logging.info(
                    "Input: %s x %s, %.2f FPS, %s frames",
                    width,
                    height,
                    fps,
                    frame_count
                )

                # Create a unique output filename.
                fd, save_path = tempfile.mkstemp(
                    suffix=".mp4",
                    prefix="animeganv3_",
                    dir="output"
                )
                os.close(fd)

                fourcc = cv2.VideoWriter_fourcc(
                    *"mp4v"
                )

                out = cv2.VideoWriter(
                    save_path,
                    fourcc,
                    fps,
                    (width, height)
                )

                if not out.isOpened():
                    raise RuntimeError(
                        "Could not initialize video writer."
                    )

                processed = 0

                # Read and process frames immediately.
                # Avoid storing batches of frames in RAM.
                while True:

                    ret, frame = cap.read()

                    if not ret:
                        break

                    stylized_frame = self.process_frame(
                        frame,
                        style_code,
                        det_face
                    )

                    # Preserve original video dimensions.
                    if (
                        stylized_frame.shape[1] != width
                        or stylized_frame.shape[0] != height
                    ):
                        stylized_frame = cv2.resize(
                            stylized_frame,
                            (width, height),
                            interpolation=cv2.INTER_LINEAR
                        )

                    out.write(stylized_frame)

                    processed += 1

                    # Log progress occasionally, not every frame.
                    if processed % 100 == 0:
                        logging.info(
                            "Processed %s / %s frames",
                            processed,
                            frame_count
                        )

                if processed == 0:
                    raise RuntimeError(
                        "No video frames were processed."
                    )

                logging.info(
                    "Conversion complete: %s frames",
                    processed
                )

                return save_path

            except gr.Error:
                raise

            except Exception as error:
                logging.exception(
                    "AnimeGANv3 conversion failed."
                )

                if save_path and os.path.exists(save_path):
                    try:
                        os.remove(save_path)
                    except OSError:
                        pass

                raise gr.Error(
                    f"Video conversion failed: {error}"
                )

            finally:

                if cap is not None:
                    cap.release()

                if out is not None:
                    out.release()


anime_gan = AnimeGANv3()


title = "AnimeGANv3: Video to Anime Converter"

description = r"""
Upload a video to convert it into anime style using AnimeGANv3.<br>
Select a style and choose whether to optimize for faces.<br>
Credits to fine tuner Asher_Chan<br>
<a href="https://github.com/TachibanaYoshino/AnimeGANv3"
target="_blank"><b>AnimeGANv3 GitHub</b></a> |
<a href="https://www.patreon.com/Asher_Chan"
target="_blank"><b>Patreon</b></a>
"""


iface = gr.Interface(
    fn=anime_gan.inference,
    inputs=[
        gr.Video(label="Input Video"),

        gr.Dropdown(
            choices=[
                "AnimeGANv3_Hayao",
                "AnimeGANv3_Shinkai",
                "AnimeGANv3_Arcane",
                "AnimeGANv3_Trump v1.0",
                "AnimeGANv3_Disney v1.0",
                "AnimeGANv3_PortraitSketch",
                "AnimeGANv3_JP_face v1.0",
                "AnimeGANv3_Kpop v2.0",
            ],
            label="AnimeGANv3 Style",
            value="AnimeGANv3_Arcane"
        ),

        gr.Radio(
            choices=["Yes", "No"],
            label="Extract face",
            value="No"
        ),
    ],

    outputs=[
        gr.Video(label="Output Video")
    ],

    title=title,
    description=description,
    allow_flagging="never"
)


if __name__ == "__main__":
    iface.queue().launch()