from __future__ import annotations

import logging

import torch

from modules import (
    devices,
    errors,
    face_restoration,
    face_restoration_utils,
    modelloader,
    shared,
)

logger = logging.getLogger(__name__)

model_url = 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/codeformer.pth'
model_download_name = 'codeformer-v0.1.0.pth'

# used by e.g. postprocessing_codeformer.py
codeformer: face_restoration.FaceRestoration | None = None


class FaceRestorerCodeFormer(face_restoration_utils.CommonFaceRestoration):
    def name(self):
        return "CodeFormer"

    def load_net(self) -> torch.Module:
        for model_path in modelloader.load_models(
            model_path=self.model_path,
            model_url=model_url,
            command_path=self.model_path,
            download_name=model_download_name,
            ext_filter=['.pth'],
        ):
            return modelloader.load_spandrel_model(
                model_path,
                device=devices.device_codeformer,
                expected_architecture='CodeFormer',
            ).model
        raise ValueError("No codeformer model found")

    def get_device(self):
        return devices.device_codeformer

    def restore(self, np_image, w: float | None = None):
        if w is None:
            w = getattr(shared.opts, "code_former_weight", 0.5)

        def restore_face(cropped_face_t):
            assert self.net is not None
            return self.net(cropped_face_t, w=w, adain=True)[0]

        return self.restore_with_helper(np_image, restore_face)


def setup_model(dirname: str) -> None:
    global codeformer
    try:
        codeformer = FaceRestorerCodeFormer(dirname)
        shared.face_restorers.append(codeformer)
    except Exception:
        errors.report("Error setting up CodeFormer", exc_info=True)