from __future__ import annotations

import io
import shutil
import sys
import threading
import urllib.request
import zipfile
from pathlib import Path

import cv2
import numpy as np
import onnxruntime as ort
import torch
import torch.nn as nn
from huggingface_hub import hf_hub_download
from insightface.app import FaceAnalysis
from torchvision.transforms.functional import rgb_to_grayscale


GHOST_CACHE = Path.home() / ".cache" / "dream-ghost2"
SOURCE_DIR = GHOST_CACHE / "source"
MODEL_REPO = "hacksider/deep-live-cam"
MODEL_REVISION = "e1c6a60039351a68150db2b50b0ef936b9ba259a"
SOURCE_ZIP = "https://github.com/ai-forever/ghost-2.0/archive/refs/heads/main.zip"
STYLEMATTE_REPO = "yc4ny/SVAD-models"
STYLEMATTE_FILE = "submodules/GAGAvatar/assets/matting/stylematte_synth.pt"

_LOCK = threading.Lock()
_ENGINE: "Ghost2Engine | None" = None


def _prepare_source() -> Path:
    marker = SOURCE_DIR / ".ready"
    if marker.exists():
        return SOURCE_DIR
    shutil.rmtree(SOURCE_DIR, ignore_errors=True)
    SOURCE_DIR.mkdir(parents=True, exist_ok=True)
    with urllib.request.urlopen(SOURCE_ZIP, timeout=120) as response:
        payload = response.read()
    with zipfile.ZipFile(io.BytesIO(payload)) as archive:
        root = archive.namelist()[0].split("/")[0]
        for member in archive.infolist():
            if not member.filename.startswith(root + "/") or member.is_dir():
                continue
            relative = Path(member.filename).relative_to(root)
            destination = SOURCE_DIR / relative
            destination.parent.mkdir(parents=True, exist_ok=True)
            with archive.open(member) as source, destination.open("wb") as target:
                shutil.copyfileobj(source, target)

    # emoca_crop is unused during inference; removing this optional training import
    # avoids pulling DECA/EMOCA/PyTorch3D into the ZeroGPU image.
    crops = SOURCE_DIR / "src" / "utils" / "crops.py"
    text = crops.read_text()
    text = text.replace(
        "from repos.emoca.gdl.datasets.ImageDatasetHelpers import bbox2point\n", ""
    )
    crops.write_text(text)
    embedder = SOURCE_DIR / "src" / "aligner" / "embedder.py"
    text = embedder.read_text()
    text = text.replace("weights='DEFAULT'", "weights=None")
    embedder.write_text(text)
    marker.write_text("ready")
    return SOURCE_DIR


def _model_file(filename: str) -> str:
    return hf_hub_download(
        repo_id=MODEL_REPO,
        filename=f"ghost2/{filename}",
        revision=MODEL_REVISION,
    )


class Ghost2Engine:
    def __init__(self) -> None:
        root = _prepare_source()
        if str(root) not in sys.path:
            sys.path.insert(0, str(root))

        from src.aligner.embedder import Embedder
        from src.aligner.generator import Generator
        from src.blender.generator import BlenderGenerator
        from src.utils.crops import norm_crop, wide_crop_face
        from src.utils.inference import copy_head_back, normalize_and_torch
        from src.utils.inpainter import LamaInpainter
        from src.utils.preblending import calc_pseudo_target_bg

        self.norm_crop = norm_crop
        self.wide_crop_face = wide_crop_face
        self.copy_head_back = copy_head_back
        self.normalize_and_torch = normalize_and_torch
        self.calc_pseudo_target_bg = calc_pseudo_target_bg

        backbone = _model_file("backbone50_1.pth")
        weights_dir = root / "weights"
        weights_dir.mkdir(exist_ok=True)
        backbone_link = weights_dir / "backbone50_1.pth"
        if not backbone_link.exists():
            backbone_link.symlink_to(backbone)

        # Upstream modules use relative asset paths.
        self._previous_cwd = Path.cwd()
        import os
        os.chdir(root)
        try:
            class AlignerInference(nn.Module):
                def __init__(inner_self) -> None:
                    super().__init__()
                    inner_self.embedder = Embedder(d_por=512, d_id=512, d_pose=256, d_exp=0)
                    inner_self.gen = Generator(
                        d_por=512, d_id=512, d_pose=256, d_exp=0,
                        padding="zero", in_channels=3, out_channels=3,
                        num_channels=64, max_num_channels=512, norm_layer="in",
                        gen_constant_input_size=4, gen_num_residual_blocks=2,
                        output_image_size=512,
                    )

                def forward(inner_self, batch):
                    return inner_self.gen(inner_self.embedder(batch))

            self.aligner = AlignerInference()
            aligner_state = torch.load(_model_file("aligner_1020_gaze_final.ckpt"), map_location="cpu")
            if "state_dict" in aligner_state:
                aligner_state = aligner_state["state_dict"]
            self.aligner.load_state_dict(
                {k: v for k, v in aligner_state.items() if k.startswith(("embedder.", "gen."))},
                strict=False,
            )

            self.blender = BlenderGenerator()
            blender_state = torch.load(_model_file("blender_lama.ckpt"), map_location="cpu")
            if "state_dict" in blender_state:
                blender_state = blender_state["state_dict"]
            self.blender.load_state_dict(
                {k.removeprefix("gen."): v for k, v in blender_state.items() if k.startswith("gen.")},
                strict=False,
            )
            self.inpainter = LamaInpainter()
        finally:
            os.chdir(self._previous_cwd)

        self.aligner = self.aligner.cuda().eval()
        self.blender = self.blender.cuda().eval()

        self.detector = FaceAnalysis(
            root=str(GHOST_CACHE / "insightface"),
            providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
            allowed_modules=["detection"],
        )
        self.detector.prepare(ctx_id=0, det_size=(640, 640))

        self.parsing = ort.InferenceSession(
            _model_file("segformer_B5_ce.onnx"),
            providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
        )
        self.parsing_input = self.parsing.get_inputs()[0].name
        self.parsing_outputs = [item.name for item in self.parsing.get_outputs()]
        self.mean = np.array([0.51315393, 0.48064056, 0.46301059])[None, :, None, None]
        self.std = np.array([0.21438347, 0.20799829, 0.20304542])[None, :, None, None]

    def _parsing(self, image: torch.Tensor) -> torch.Tensor:
        prepared = (((image[:, [2, 1, 0]] / 2 + 0.5).detach().cpu().numpy() - self.mean) / self.std)
        result = self.parsing.run(
            self.parsing_outputs, {self.parsing_input: prepared.astype(np.float32)}
        )[0]
        return torch.tensor(result, device="cuda", dtype=torch.float32)

    @staticmethod
    def _head_mask(parsing: torch.Tensor) -> torch.Tensor:
        mask = torch.zeros_like(parsing, dtype=torch.float32)
        for index in range(1, 21):
            mask[parsing == index] = 1.0
        return mask[0, 0] if mask.ndim == 4 else mask[0]

    def _process(self, frame: np.ndarray, target: bool = False):
        faces = self.detector.get(frame)
        if not faces:
            raise ValueError("No head was detected in the image/frame.")
        face = max(
            faces,
            key=lambda item: float((item.bbox[2] - item.bbox[0]) * (item.bbox[3] - item.bbox[1])),
        )
        keypoints = face.kps
        wide = self.wide_crop_face(frame, keypoints, return_M=target)
        if target:
            wide, matrix = wide
        arc = self.norm_crop(frame, keypoints)
        arc_tensor = self.normalize_and_torch(arc)
        wide_tensor = self.normalize_and_torch(wide)
        mask = self._head_mask(self._parsing(wide_tensor))
        if target:
            return wide_tensor, arc_tensor, mask, frame, matrix
        return wide_tensor, arc_tensor, mask

    def prepare_source(self, source_bgr: np.ndarray) -> dict[str, torch.Tensor]:
        wide, arc, mask = self._process(source_bgr)
        return {
            "wide": wide.unsqueeze(1),
            "arc": arc.unsqueeze(1),
            "mask": mask,
        }

    def swap_frame(self, source: dict[str, torch.Tensor], target_bgr: np.ndarray) -> np.ndarray:
        wide_target, arc_target, target_mask, full_frame, matrix = self._process(
            target_bgr, target=True
        )
        source_mask = source["mask"]
        batch = {
            "source": {
                "face_arc": source["arc"],
                "face_wide": source["wide"] * source_mask,
                "face_wide_mask": source_mask,
            },
            "target": {
                "face_arc": arc_target,
                "face_wide": wide_target * target_mask,
                "face_wide_mask": target_mask,
            },
        }
        with torch.inference_mode():
            aligned = self.aligner(batch)
            target_parsing = self._parsing(wide_target)
            pseudo_background = self.calc_pseudo_target_bg(wide_target, target_parsing)
            aligned_parsing = self._parsing(aligned["fake_rgbs"] * aligned["fake_segm"])
            soft_mask = self._head_mask(aligned_parsing).unsqueeze(0)
            new_source = (
                aligned["fake_rgbs"] * soft_mask[:, None]
                + pseudo_background * (1 - soft_mask[:, None])
            )
            output = self.blender(
                new_source,
                rgb_to_grayscale(new_source[0][[2, 1, 0]]).unsqueeze(0),
                wide_target,
                aligned_parsing,
                target_parsing,
                gt=wide_target,
                M_a_noise=None,
                M_t_noise=None,
                cycle=False,
                train=False,
                return_inputs=True,
                inpainter=self.inpainter,
            )[0]
        crop = np.uint8(
            (output[0].detach().cpu().numpy().transpose(1, 2, 0)[:, :, ::-1] / 2 + 0.5) * 255
        )
        rgb = self.copy_head_back(crop, full_frame[..., ::-1], matrix)
        return rgb[..., ::-1].copy()


def get_ghost_engine() -> Ghost2Engine:
    global _ENGINE
    with _LOCK:
        if _ENGINE is None:
            _ENGINE = Ghost2Engine()
        return _ENGINE