#!/usr/bin/env python3
"""Create the exact opposite direction of an extracted ComfyUI LoRA.

For a normal LoRA patch, negating either the up or down factor negates the
represented matrix delta. ComfyUI ``.diff`` and ``.diff_b`` tensors are
negated directly. Alpha and the other LoRA factor remain unchanged.

The conversion streams one tensor at a time and writes resumable safetensors
parts, so even very large LoRAs require little memory.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import shutil
from pathlib import Path
from typing import Any

import torch
from safetensors import safe_open
from safetensors.torch import save_file

from extract_minimax_h3_ref2va_lora import merge_parts, write_json_atomic


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Invert an extracted ComfyUI LoRA without recomputing its SVD."
    )
    parser.add_argument("--input", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--keep-parts", action="store_true")
    parser.add_argument("--quiet", action="store_true")
    return parser.parse_args()


def should_negate(key: str) -> bool:
    return (
        key.endswith(".lora_up.weight")
        or key.endswith(".diff")
        or key.endswith(".diff_b")
    )


def part_path(parts_dir: Path, index: int, key: str) -> Path:
    digest = hashlib.sha1(key.encode("utf-8")).hexdigest()[:12]
    return parts_dir / f"{index:04d}-{digest}.safetensors"


def reverse_metadata(metadata: dict[str, str], input_path: Path) -> dict[str, str]:
    output = dict(metadata)
    source_base = output.get("source_base")
    source_target = output.get("source_target")
    if source_base is not None and source_target is not None:
        output["source_base"] = source_target
        output["source_target"] = source_base
    output["modelspec.title"] = "MiniMax-H3 REF2VA to FL2VA extracted LoRA"
    output["modelspec.description"] = (
        "Opposite-direction form of an extracted FL2VA-to-REF2VA LoRA; "
        "one factor and all exact ComfyUI diff patches are sign-inverted."
    )
    output["direction"] = "ref2va_to_fl2va"
    output["inverted_from"] = input_path.name
    return output


def reverse_report(input_path: Path, output_path: Path, summary: dict[str, Any]) -> None:
    input_report = input_path.with_suffix(input_path.suffix + ".json")
    output_report = output_path.with_suffix(output_path.suffix + ".json")
    if input_report.is_file():
        report = json.loads(input_report.read_text())
        report_summary = report.setdefault("summary", {})
        report_summary.update(summary)
        write_json_atomic(output_report, report)
    else:
        write_json_atomic(output_report, {"summary": summary})


def run(args: argparse.Namespace) -> None:
    if not args.input.is_file():
        raise FileNotFoundError(args.input)
    if args.output.exists():
        raise FileExistsError(f"Refusing to overwrite {args.output}")

    parts_dir = Path(f"{args.output}.parts")
    parts_dir.mkdir(parents=True, exist_ok=True)
    progress_path = parts_dir / "progress.json"
    progress: dict[str, Any] = {
        "input": str(args.input.resolve()),
        "completed": {},
    }
    if progress_path.exists():
        progress = json.loads(progress_path.read_text())
        if progress.get("input") != str(args.input.resolve()):
            raise ValueError("Existing parts belong to a different input file")
    completed: dict[str, Any] = progress.setdefault("completed", {})

    with safe_open(args.input, framework="pt", device="cpu") as source:
        keys = list(source.keys())
        metadata = source.metadata() or {}
        negated_count = 0
        copied_count = 0

        for index, key in enumerate(keys):
            output_part = part_path(parts_dir, index, key)
            negate = should_negate(key)
            previous = completed.get(key)
            if previous is not None and output_part.is_file():
                if previous["negated"]:
                    negated_count += 1
                else:
                    copied_count += 1
                if not args.quiet:
                    print(f"[{index + 1:04d}/{len(keys)}] resume {key}", flush=True)
                continue

            tensor = source.get_tensor(key)
            output_tensor = (-tensor).contiguous() if negate else tensor.contiguous()
            temporary = output_part.with_name(f".{output_part.name}.tmp")
            temporary.unlink(missing_ok=True)
            save_file({key: output_tensor}, str(temporary))
            os.replace(temporary, output_part)
            completed[key] = {
                "negated": negate,
                "shape": list(tensor.shape),
                "dtype": str(tensor.dtype),
                "part": output_part.name,
            }
            write_json_atomic(progress_path, progress)
            if negate:
                negated_count += 1
                action = "negate"
            else:
                copied_count += 1
                action = "copy"
            if not args.quiet:
                print(f"[{index + 1:04d}/{len(keys)}] {action} {key}", flush=True)

    parts = sorted(parts_dir.glob("[0-9][0-9][0-9][0-9]-*.safetensors"))
    output_metadata = reverse_metadata(metadata, args.input)
    tensor_count = merge_parts(parts, args.output, output_metadata)
    if tensor_count != len(keys):
        raise ValueError(f"Final tensor count mismatch: {tensor_count} != {len(keys)}")

    with safe_open(args.output, framework="pt", device="cpu") as final:
        if list(final.keys()) != keys:
            raise ValueError("Final key order/content differs from the input")

    summary = {
        "direction": "ref2va_to_fl2va",
        "inverted_from": args.input.name,
        "output_file": args.output.name,
        "output_tensors": tensor_count,
        "negated_tensors": negated_count,
        "copied_tensors": copied_count,
        "output_bytes": args.output.stat().st_size,
    }
    reverse_report(args.input, args.output, summary)
    print(json.dumps(summary, indent=2), flush=True)

    if not args.keep_parts:
        for part in parts:
            part.unlink()
        progress_path.unlink(missing_ok=True)
        shutil.rmtree(parts_dir)


def main() -> None:
    run(parse_args())


if __name__ == "__main__":
    main()