"""One resident pipeline; original backbone tensors are never cumulatively merged."""
from pathlib import Path
import torch

from anima_role_ref_kv.components import load_transformer, PromptEncoder, VAE, NEGATIVE
from anima_role_ref_kv.pipeline import sample, resize_reference
from anima_role_ref_kv.weights import load, read_state, swap
from checkpoints import native_plan, diffusers_plan, copy_plan


class Runtime:
    def __init__(self, base_dir, adapter_file, device="cuda", adapter_key="v2.0"):
        self.base_dir = str(base_dir)
        self.net = load_transformer(base_dir, device)
        self.text = PromptEncoder(base_dir, device)
        self.vae = VAE(base_dir, device)
        # Capture references before LoRALinear wraps linears. The adapter keeps
        # its frozen base outside its registered parameters; net.state_dict()
        # after injection would omit those base weights.
        self.targets = {("transformer", k): v for k, v in self.net.state_dict(keep_vars=True).items()}
        self.targets.update({("conditioner", k): v for k, v in
                             self.text.components.text_conditioner.state_dict(keep_vars=True).items()})
        self.base_plan = diffusers_plan(base_dir, self.targets)
        self.adapter_states = {adapter_key: read_state(self.net, adapter_file)}
        self.adapter = load(self.net, adapter_file)
        self.current = "base-v1.0"
        self.current_adapter = adapter_key

    def select(self, model_key, filename=None):
        if self.current == model_key:
            return
        plan = (self.base_plan if model_key == "base-v1.0" else
                (diffusers_plan(filename, self.targets) if Path(filename).is_dir() else
                 native_plan(filename, self.targets)))
        self.current = None  # A failed copy must force a complete reload next time.
        copy_plan(plan, self.targets)
        self.current = model_key

    def select_adapter(self, adapter_key, filename=None):
        if self.current_adapter == adapter_key:
            return
        state = self.adapter_states.get(adapter_key)
        if state is None:
            if not filename:
                raise ValueError(f"Adapter file required for {adapter_key}")
            state = read_state(self.net, filename)
            self.adapter_states[adapter_key] = state
        self.current_adapter = None  # A failed copy must force a complete reload next time.
        swap(self.adapter, state)
        self.current_adapter = adapter_key

    @torch.no_grad()
    def generate(self, model_key, filename, reference, prompt, *, adapter_key=None,
                 adapter_file=None, strength=1.0, negative_prompt=NEGATIVE,
                 progress_callback=None, **settings):
        self.select(model_key, filename)
        if adapter_key is not None:
            self.select_adapter(adapter_key, adapter_file)
        for layer in self.adapter.layers.values():
            layer.multiplier = strength
        cfg = settings.get("cfg", 4.)
        positive = self.text([prompt])
        negative = self.text([negative_prompt]) if cfg != 1 else None
        ref = self.vae.encode(resize_reference(reference))
        latent = sample(self.net, positive, negative, ref,
                        progress_callback=progress_callback, **settings)
        if not torch.isfinite(latent).all():
            raise FloatingPointError("The model produced non-finite latents")
        return self.vae.decode(latent)