import os
import argparse

import cv2
import numpy as np
import torch
from safetensors.torch import load_file
from transformers import AutoImageProcessor

from lib.utils_segfly import ID2COLOR, mask2label
from lib.firefly_rgb import FireflyForSemanticSegmentationRGB, FireflyConfigRGB
from lib.firefly_thermal import FireflyForSemanticSegmentationThermal, FireflyConfigThermal


def run_infer(args):
    if not args.weights_path:
        if args.modality == "rgb":
            args.weights_path = "./Firefly_RGB/model.safetensors"
        else:
            args.weights_path = "./Firefly_Thermal/model.safetensors"
        print(f"No weights path provided, defaulting to: {args.weights_path}")

    weights_dir = os.path.dirname(args.weights_path) if os.path.isfile(args.weights_path) else args.weights_path

    print(f"Loading image processor from {weights_dir}...")
    image_processor = AutoImageProcessor.from_pretrained(weights_dir)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    print(f"Initializing {args.modality} model...")
    if args.modality == "rgb":
        cfg = FireflyConfigRGB(
            num_labels=15,
            image_size=args.image_size,
            embedding_dim=256,
            backbone_embed_dim=768,
            patch_size=16,
            repo_dir=args.dinov3_repo_dir,
            model_name="dinov3_vitb16",
            semantic_loss_ignore_index=255,
        )
        model = FireflyForSemanticSegmentationRGB(cfg)
    else:
        cfg = FireflyConfigThermal(
            num_labels=15,
            image_size=args.image_size,
            embedding_dim=256,
            backbone_embed_dim=768,
            patch_size=16,
            num_layers=12,
            rein_token_length=100,
            feature_layers=[2, 5, 8, 11],
            repo_dir=args.dinov3_repo_dir,
            model_name="dinov3_vitb16",
            semantic_loss_ignore_index=255,
        )
        model = FireflyForSemanticSegmentationThermal(cfg)

    print(f"Loading weights from {args.weights_path}...")
    state_dict = load_file(args.weights_path)
    new_state_dict = {
        k[7:] if k.startswith("module.") else k: v
        for k, v in state_dict.items()
    }
    msg = model.load_state_dict(new_state_dict, strict=False)
    print(f"Weights loaded. Missing: {len(msg.missing_keys)}, Unexpected: {len(msg.unexpected_keys)}")

    model.eval()
    model.to(device)

    image_bgr = cv2.imread(args.image)
    if image_bgr is None:
        raise FileNotFoundError(f"Could not read image: {args.image}")
    image_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
    orig_h, orig_w = image_rgb.shape[:2]

    inputs = image_processor(images=image_rgb, return_tensors="pt")
    pixel_values = inputs.pixel_values.to(device, dtype=torch.float32)

    with torch.no_grad():
        logits = model(pixel_values).logits
        upsampled = torch.nn.functional.interpolate(
            logits.float(), size=(orig_h, orig_w), mode="bilinear", align_corners=False
        )
        pred = upsampled.argmax(dim=1).squeeze(0).cpu().numpy()

    os.makedirs(args.output, exist_ok=True)

    pred_color = mask2label(pred, ID2COLOR)
    pred_bgr = cv2.cvtColor(pred_color, cv2.COLOR_RGB2BGR)

    stem = os.path.splitext(os.path.basename(args.image))[0]
    out_path = os.path.join(args.output, f"{stem}_pred.png")
    cv2.imwrite(out_path, pred_bgr)
    print(f"Segmentation saved to {out_path}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Single-image inference for Firefly RGB / Thermal models")
    parser.add_argument("--image", type=str, required=True, help="Path to the input image.")
    parser.add_argument("--modality", type=str, default="rgb", choices=["rgb", "thermal"],
                        help="Model modality: 'rgb' or 'thermal'.")
    parser.add_argument("--output", type=str, default="./infer_output",
                        help="Directory to save the colorized segmentation output.")
    parser.add_argument("--weights_path", type=str, default="",
                        help="Path to model weights (.safetensors). Auto-detected from modality if not set.")
    parser.add_argument("--dinov3_repo_dir", type=str, default="./dinov3",
                        help="Path to the local DINOv3 repository.")
    parser.add_argument("--image_size", type=int, default=640,
                        help="Resolution at which the image is fed to the model.")
    args = parser.parse_args()
    run_infer(args)