from argparse import ArgumentParser
from pathlib import Path
from typing import Iterable, Dict, List

from pydantic_xml import BaseXmlModel, attr, element

# --- Pydantic XML models ---

class CVATLabel(BaseXmlModel, tag="label"):
    name: str = element()

class CVATTask(BaseXmlModel, tag="task"):
    labels: list[CVATLabel] = element(tag="label")

class CVATJob(BaseXmlModel, tag="job"):
    labels: list[CVATLabel] = element(tag="label")

class CVATProject(BaseXmlModel, tag="project"):
    labels: list[CVATLabel] = element(tag="label")

class CVATMeta(BaseXmlModel, tag="meta"):
    task: CVATTask | None = element(default=None)
    job: CVATJob | None = element(default=None)
    project: CVATProject | None = element(default=None)

class CVATBox(BaseXmlModel, tag="box"):
    label: str = attr()
    xtl: float = attr()
    ytl: float = attr()
    xbr: float = attr()
    ybr: float = attr()

class CVATImage(BaseXmlModel, tag="image"):
    id: int = attr()
    name: str = attr()  # Filename
    width: int = attr()
    height: int = attr()
    boxes: list[CVATBox] = element(tag="box", default=[])

class CVATAnnotations(BaseXmlModel, tag="annotations"):
    meta: CVATMeta = element()
    images: list[CVATImage] = element(tag="image", default=[])

# --- Input classes ---

class _CVATBaseInput:
    @staticmethod
    def add_cli_arguments(parser: ArgumentParser) -> None:
        parser.add_argument(
            "--input-file",
            type=Path,
            required=True,
            help="Path to input CVAT XML file",
        )

    def __init__(self, input_file: Path) -> None:
        with input_file.open("r", encoding="utf-8") as file:
            xml_content = file.read()
        self._data = CVATAnnotations.from_xml(xml_content)

    def get_categories(self) -> Iterable["Category"]:
        meta = self._data.meta
        labels: List[CVATLabel] | None = None
        if meta.task is not None and meta.task.labels:
            labels = meta.task.labels
        elif meta.job is not None and meta.job.labels:
            labels = meta.job.labels
        elif meta.project is not None and meta.project.labels:
            labels = meta.project.labels
        if labels is None:
            raise ValueError(
                "Could not find labels in meta/task, meta/job, or meta/project"
            )
        for idx, label in enumerate(labels, start=1):
            yield Category(id=idx, name=label.name)

    def get_images(self) -> Iterable["Image"]:
        for img in self._data.images:
            yield Image(
                id=img.id,
                filename=img.name,
                width=img.width,
                height=img.height,
            )

class CVATObjectDetectionInput(_CVATBaseInput, ObjectDetectionInput):
    def get_labels(self) -> Iterable["ImageObjectDetection"]:
        category_by_name: Dict[str, Category] = {
            cat.name: cat for cat in self.get_categories()
        }
        for img in self._data.images:
            objects = []
            for box in img.boxes:
                cat = category_by_name.get(box.label)
                if cat is None:
                    continue  # or handle error as needed
                objects.append(
                    SingleObjectDetection(
                        category=cat,
                        box=BoundingBox.from_format(
                            bbox=[box.xtl, box.ytl, box.xbr, box.ybr],
                            format=BoundingBoxFormat.XYXY,
                        ),
                    )
                )
            yield ImageObjectDetection(
                image=Image(id=img.id, filename=img.name, width=img.width, height=img.height),
                objects=objects,
            )