"""Run the 2D floor-plan component detector from Hugging Face Hub or local files.

Example:
    python inference.py plan.png --output-dir results/plan
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np
import torch
from huggingface_hub import hf_hub_download
from PIL import Image, ImageDraw
from rfdetr import RFDETRMedium


DEFAULT_REPO_ID = "OsamaMo/2dplan2strct"
COLORS = {
    "wall": "#e63946",
    "room": "#457b9d",
    "door": "#f4a261",
    "window": "#2a9d8f",
}


class FloorPlanDetector:
    """Load RF-DETR once, then run it on one or more floor-plan images."""

    def __init__(
        self,
        repo_id: str = DEFAULT_REPO_ID,
        *,
        revision: str | None = None,
        model_dir: Path | None = None,
        device: str = "auto",
    ) -> None:
        self.repo_id = repo_id
        if device not in {"auto", "cpu", "cuda"}:
            raise ValueError("device must be 'auto', 'cpu', or 'cuda'")
        self.device = "cuda" if device == "auto" and torch.cuda.is_available() else device
        if self.device == "auto":
            self.device = "cpu"
        if self.device == "cuda" and not torch.cuda.is_available():
            raise RuntimeError("CUDA was requested, but no CUDA GPU is available")
        if model_dir is None:
            config_path = Path(hf_hub_download(repo_id=repo_id, filename="config.json", revision=revision))
            config = json.loads(config_path.read_text(encoding="utf-8"))
            weights = hf_hub_download(repo_id=repo_id, filename=config["checkpoint"], revision=revision)
        else:
            model_dir = Path(model_dir)
            config = json.loads((model_dir / "config.json").read_text(encoding="utf-8"))
            weights = str(model_dir / config["checkpoint"])

        if config["variant"] != "RFDETRMedium":
            raise ValueError(f"Unsupported model variant: {config['variant']}")
        self.class_names = config["class_names"]
        self.model = RFDETRMedium(
            pretrain_weights=weights,
            resolution=config["resolution"],
            num_classes=config["num_classes"],
            device=self.device,
        )
        if self.model.class_names and self.model.class_names != self.class_names:
            raise ValueError("Checkpoint class names do not match config.json")

    def predict(self, image_path: str | Path, *, threshold: float = 0.35) -> dict:
        """Return boxes in original-image pixel coordinates, sorted by confidence."""
        if not 0 <= threshold <= 1:
            raise ValueError("threshold must be between 0 and 1")
        image_path = Path(image_path)
        with Image.open(image_path) as source:
            image = np.array(source.convert("RGB"))
        height, width = image.shape[:2]
        result = self.model.predict(image, threshold=threshold)
        detections = [
            {
                "label": self.class_names[int(class_id)],
                "score": float(score),
                "box_xyxy": [float(value) for value in box],
            }
            for box, score, class_id in zip(result.xyxy, result.confidence, result.class_id)
        ]
        detections.sort(key=lambda item: item["score"], reverse=True)
        return {
            "model": self.repo_id,
            "image": str(image_path),
            "width": width,
            "height": height,
            "threshold": threshold,
            "detections": detections,
        }


def save_prediction(prediction: dict, image_path: str | Path, output_dir: str | Path) -> tuple[Path, Path]:
    """Write a JSON result and an annotated PNG without changing the source image."""
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    json_path = output_dir / "predictions.json"
    image_out = output_dir / "annotated.png"
    json_path.write_text(json.dumps(prediction, indent=2) + "\n", encoding="utf-8")

    with Image.open(image_path) as source:
        image = source.convert("RGB")
    draw = ImageDraw.Draw(image)
    line_width = max(2, min(image.size) // 300)
    for item in reversed(prediction["detections"]):
        color = COLORS.get(item["label"], "#ffffff")
        box = item["box_xyxy"]
        draw.rectangle(box, outline=color, width=line_width)
        label = f"{item['label']} {item['score']:.2f}"
        text_box = draw.textbbox((box[0], box[1]), label)
        text_height = text_box[3] - text_box[1]
        label_y = max(0, box[1] - text_height - 4)
        draw.rectangle((box[0], label_y, box[0] + text_box[2] - text_box[0] + 4, label_y + text_height + 4), fill=color)
        draw.text((box[0] + 2, label_y + 2), label, fill="white")
    image.save(image_out)
    return json_path, image_out


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("image", type=Path, help="PNG or JPEG floor-plan image")
    parser.add_argument("--repo-id", default=DEFAULT_REPO_ID, help="Hugging Face model repository")
    parser.add_argument("--revision", help="Optional Hub commit SHA or tag for reproducible downloads")
    parser.add_argument("--model-dir", type=Path, help="Use previously downloaded config and weights offline")
    parser.add_argument("--device", choices=("auto", "cpu", "cuda"), default="auto",
                        help="Run on CPU or CUDA; auto uses CUDA when available")
    parser.add_argument("--threshold", type=float, default=0.35, help="Confidence threshold from 0 to 1")
    parser.add_argument("--output-dir", type=Path, default=Path("output/prediction"))
    args = parser.parse_args()
    if not args.image.is_file():
        parser.error(f"image does not exist: {args.image}")
    if not 0 <= args.threshold <= 1:
        parser.error("--threshold must be between 0 and 1")

    detector = FloorPlanDetector(args.repo_id, revision=args.revision, model_dir=args.model_dir,
                                 device=args.device)
    prediction = detector.predict(args.image, threshold=args.threshold)
    json_path, image_path = save_prediction(prediction, args.image, args.output_dir)
    print(f"{len(prediction['detections'])} detections")
    print(f"JSON: {json_path}")
    print(f"Annotated image: {image_path}")


if __name__ == "__main__":
    main()