from dataclasses import dataclass
from functools import lru_cache
from inspect import signature
from pathlib import Path

import gradio as gr
import numpy as np
from huggingface_hub import hf_hub_download
from PIL import Image, ImageDraw
from ultralytics import YOLO


APP_NAME = "AI-PPE-Detection-System"
PPE_MODEL_REPO = "Hexmon/vyra-yolo-ppe-detection"
PPE_MODEL_FILE = "best.pt"
PERSON_MODEL_NAME = "yolov8n.pt"

PERSON_CONFIDENCE_THRESHOLD = 0.40
PPE_CONFIDENCE_THRESHOLD = 0.50
IOU_THRESHOLD = 0.45
IMAGE_SIZE = 960
PPE_PERSON_OVERLAP_THRESHOLD = 0.20

HELMET_LABELS = {"helmet", "hardhat", "hard hat", "safety helmet", "with helmet"}
VEST_LABELS = {"vest", "safety vest", "with vest", "reflective vest", "safety jacket"}
NO_HELMET_LABELS = {"no helmet", "no hardhat", "no hard hat", "without helmet"}
NO_VEST_LABELS = {"no vest", "no safety vest", "without vest"}
SUPPORTED_PPE_LABELS = HELMET_LABELS | VEST_LABELS | NO_HELMET_LABELS | NO_VEST_LABELS

STATUS_OK = "OK"
STATUS_NG = "NG"
STATUS_UNKNOWN = "Unknown"


@dataclass(frozen=True)
class DetectionBox:
    x1: float
    y1: float
    x2: float
    y2: float
    confidence: float
    label: str
    category: str

    @property
    def area(self) -> float:
        return max(0.0, self.x2 - self.x1) * max(0.0, self.y2 - self.y1)

    @property
    def center(self) -> tuple[float, float]:
        return (self.x1 + self.x2) / 2, (self.y1 + self.y2) / 2


@dataclass(frozen=True)
class WorkerAssessment:
    worker_id: int
    worker_box: DetectionBox
    helmet_status: str
    vest_status: str
    helmet_confidence: float | None
    vest_confidence: float | None

    @property
    def overall_status(self) -> str:
        return STATUS_OK if self.helmet_status == STATUS_OK and self.vest_status == STATUS_OK else STATUS_NG

    @property
    def missing_items(self) -> str:
        missing = []
        if self.helmet_status != STATUS_OK:
            missing.append("Helmet")
        if self.vest_status != STATUS_OK:
            missing.append("Safety Vest")
        return ", ".join(missing) if missing else "None"


@lru_cache(maxsize=1)
def load_ppe_model() -> YOLO:
    model_path = hf_hub_download(repo_id=PPE_MODEL_REPO, filename=PPE_MODEL_FILE)
    return YOLO(model_path)


@lru_cache(maxsize=1)
def load_person_model() -> YOLO:
    return YOLO(PERSON_MODEL_NAME)


def normalize_label(label: str) -> str:
    return label.strip().lower().replace("_", " ").replace("-", " ")


def classify_ppe_label(label: str) -> str:
    normalized = normalize_label(label)
    if normalized in HELMET_LABELS:
        return "helmet"
    if normalized in VEST_LABELS:
        return "vest"
    if normalized in NO_HELMET_LABELS:
        return "no_helmet"
    if normalized in NO_VEST_LABELS:
        return "no_vest"
    return "other"


def intersection_area(first: DetectionBox, second: DetectionBox) -> float:
    x1 = max(first.x1, second.x1)
    y1 = max(first.y1, second.y1)
    x2 = min(first.x2, second.x2)
    y2 = min(first.y2, second.y2)
    return max(0.0, x2 - x1) * max(0.0, y2 - y1)


def point_in_box(point: tuple[float, float], box: DetectionBox, margin_ratio: float = 0.08) -> bool:
    center_x, center_y = point
    width = box.x2 - box.x1
    height = box.y2 - box.y1
    margin_x = width * margin_ratio
    margin_y = height * margin_ratio

    return (
        box.x1 - margin_x <= center_x <= box.x2 + margin_x
        and box.y1 - margin_y <= center_y <= box.y2 + margin_y
    )


def person_association_score(ppe_box: DetectionBox, worker_box: DetectionBox) -> float:
    overlap_ratio = intersection_area(ppe_box, worker_box) / max(ppe_box.area, 1.0)
    center_bonus = 1.0 if point_in_box(ppe_box.center, worker_box) else 0.0
    return overlap_ratio + center_bonus


def assign_ppe_to_workers(
    workers: list[DetectionBox],
    ppe_boxes: list[DetectionBox],
) -> dict[int, list[DetectionBox]]:
    assignments = {index: [] for index in range(len(workers))}

    for ppe_box in ppe_boxes:
        scored_workers = [
            (index, person_association_score(ppe_box, worker_box))
            for index, worker_box in enumerate(workers)
        ]
        if not scored_workers:
            continue

        best_worker_index, best_score = max(scored_workers, key=lambda item: item[1])
        overlap_ratio = intersection_area(ppe_box, workers[best_worker_index]) / max(ppe_box.area, 1.0)
        if best_score >= 1.0 or overlap_ratio >= PPE_PERSON_OVERLAP_THRESHOLD:
            assignments[best_worker_index].append(ppe_box)

    return assignments


def build_ppe_assignment_lookup(worker_ppe: dict[int, list[DetectionBox]]) -> dict[DetectionBox, int]:
    return {
        ppe_box: worker_index
        for worker_index, assigned_boxes in worker_ppe.items()
        for ppe_box in assigned_boxes
    }


def pick_best_confidence(boxes: list[DetectionBox], categories: set[str]) -> float | None:
    candidates = [box.confidence for box in boxes if box.category in categories]
    return max(candidates) if candidates else None


def assess_workers(
    workers: list[DetectionBox],
    worker_ppe: dict[int, list[DetectionBox]],
) -> list[WorkerAssessment]:
    assessments = []
    for index, worker_box in enumerate(workers):
        assigned_boxes = worker_ppe.get(index, [])

        helmet_confidence = pick_best_confidence(assigned_boxes, {"helmet"})
        vest_confidence = pick_best_confidence(assigned_boxes, {"vest"})
        no_helmet_confidence = pick_best_confidence(assigned_boxes, {"no_helmet"})
        no_vest_confidence = pick_best_confidence(assigned_boxes, {"no_vest"})

        helmet_status = STATUS_OK if helmet_confidence is not None else STATUS_NG
        vest_status = STATUS_OK if vest_confidence is not None else STATUS_NG

        if no_helmet_confidence is not None and (helmet_confidence is None or no_helmet_confidence > helmet_confidence):
            helmet_status = STATUS_NG
        if no_vest_confidence is not None and (vest_confidence is None or no_vest_confidence > vest_confidence):
            vest_status = STATUS_NG

        assessments.append(
            WorkerAssessment(
                worker_id=index + 1,
                worker_box=worker_box,
                helmet_status=helmet_status,
                vest_status=vest_status,
                helmet_confidence=helmet_confidence,
                vest_confidence=vest_confidence,
            )
        )

    return assessments


def extract_person_boxes(person_result) -> list[DetectionBox]:
    boxes = []
    for box in person_result.boxes:
        x1, y1, x2, y2 = box.xyxy[0].tolist()
        boxes.append(DetectionBox(x1, y1, x2, y2, float(box.conf[0]), "person", "worker"))
    return boxes


def extract_ppe_boxes(ppe_result) -> list[DetectionBox]:
    boxes = []
    for box in ppe_result.boxes:
        class_id = int(box.cls[0])
        label = normalize_label(ppe_result.names[class_id])
        if label not in SUPPORTED_PPE_LABELS:
            continue
        x1, y1, x2, y2 = box.xyxy[0].tolist()
        boxes.append(DetectionBox(x1, y1, x2, y2, float(box.conf[0]), label, classify_ppe_label(label)))
    return boxes


def format_confidence(confidence: float | None) -> str:
    return f"{confidence:.2f}" if confidence is not None else "-"


def build_ppe_markdown(assessments: list[WorkerAssessment], ppe_boxes: list[DetectionBox]) -> str:
    detected_helmet = any(box.category == "helmet" for box in ppe_boxes)
    detected_vest = any(box.category == "vest" for box in ppe_boxes)
    all_workers_ok = bool(assessments) and all(worker.overall_status == STATUS_OK for worker in assessments)

    return f"""
### PPE Check Result

| Item | Status |
|---|---|
| Helmet | {STATUS_OK if detected_helmet else STATUS_NG} |
| Safety Vest | {STATUS_OK if detected_vest else STATUS_NG} |
| Overall Result | {STATUS_OK if all_workers_ok else STATUS_NG} |
"""


def build_worker_summary_markdown(assessments: list[WorkerAssessment]) -> str:
    worker_count = len(assessments)
    if worker_count == 0:
        return """
### Worker Safety Summary

| Metric | Value |
|---|---|
| Workers | 0 |
| Fully Compliant Workers | No workers detected |
| Helmet Compliance | No workers detected |
| Vest Compliance | No workers detected |
"""

    helmet_ok_count = sum(worker.helmet_status == STATUS_OK for worker in assessments)
    vest_ok_count = sum(worker.vest_status == STATUS_OK for worker in assessments)
    fully_compliant_count = sum(worker.overall_status == STATUS_OK for worker in assessments)

    return f"""
### Worker Safety Summary

| Metric | Value |
|---|---|
| Workers | {worker_count} |
| Fully Compliant Workers | {fully_compliant_count}/{worker_count} |
| Helmet Compliance | {helmet_ok_count}/{worker_count} |
| Vest Compliance | {vest_ok_count}/{worker_count} |
"""


def build_status_markdown(assessments: list[WorkerAssessment]) -> str:
    if not assessments:
        return """
### Site Safety Status

**NG - No workers detected**

Upload a clearer image or lower the worker confidence threshold.
"""

    non_compliant = [worker for worker in assessments if worker.overall_status == STATUS_NG]
    if not non_compliant:
        return """
### Site Safety Status

**OK - All detected workers are PPE compliant**

Every detected worker has both helmet and safety vest evidence.
"""

    missing_summary = "; ".join(
        f"Worker {worker.worker_id}: {worker.missing_items}" for worker in non_compliant
    )
    return f"""
### Site Safety Status

**NG - {len(non_compliant)} worker(s) need attention**

{missing_summary}
"""


def build_worker_table(assessments: list[WorkerAssessment]) -> list[list[str]]:
    return [
        [
            f"Worker {worker.worker_id}",
            worker.helmet_status,
            worker.vest_status,
            worker.overall_status,
            worker.missing_items,
            format_confidence(worker.worker_box.confidence),
            format_confidence(worker.helmet_confidence),
            format_confidence(worker.vest_confidence),
        ]
        for worker in assessments
    ]


def build_detection_table(
    workers: list[DetectionBox],
    ppe_boxes: list[DetectionBox],
    ppe_assignment: dict[DetectionBox, int],
) -> list[list[str]]:
    rows = [
        [
            "worker",
            box.label,
            format_confidence(box.confidence),
            "-",
            f"{int(box.x1)}, {int(box.y1)}, {int(box.x2)}, {int(box.y2)}",
        ]
        for box in workers
    ]
    rows.extend(
        [
            box.category,
            box.label,
            format_confidence(box.confidence),
            f"Worker {ppe_assignment[box] + 1}" if box in ppe_assignment else "Unassigned",
            f"{int(box.x1)}, {int(box.y1)}, {int(box.x2)}, {int(box.y2)}",
        ]
        for box in ppe_boxes
    )
    return rows


def draw_label(draw: ImageDraw.ImageDraw, x: int, y: int, label: str, fill: tuple[int, int, int]) -> None:
    text_bbox = draw.textbbox((x, y), label)
    text_width = text_bbox[2] - text_bbox[0]
    text_height = text_bbox[3] - text_bbox[1]
    label_y = max(0, y - text_height - 8)
    draw.rectangle((x, label_y, x + text_width + 8, y), fill=fill)
    draw.text((x + 4, label_y + 2), label, fill=(255, 255, 255))


def draw_detection_boxes(
    image: Image.Image,
    workers: list[DetectionBox],
    ppe_boxes: list[DetectionBox],
    assessments: list[WorkerAssessment],
    ppe_assignment: dict[DetectionBox, int],
) -> Image.Image:
    annotated_image = image.copy()
    draw = ImageDraw.Draw(annotated_image)
    assessment_by_id = {assessment.worker_id: assessment for assessment in assessments}

    for index, worker_box in enumerate(workers, start=1):
        assessment = assessment_by_id.get(index)
        is_ok = assessment is not None and assessment.overall_status == STATUS_OK
        color = (30, 150, 85) if is_ok else (220, 80, 60)
        x1, y1, x2, y2 = map(int, (worker_box.x1, worker_box.y1, worker_box.x2, worker_box.y2))
        draw.rectangle((x1, y1, x2, y2), outline=color, width=4)
        status = assessment.overall_status if assessment else STATUS_UNKNOWN
        draw_label(draw, x1, y1, f"Worker {index}: {status}", color)

    for box in ppe_boxes:
        if box not in ppe_assignment:
            color = (100, 116, 139)
        elif box.category in {"helmet", "vest"}:
            color = (35, 160, 80)
        elif box.category in {"no_helmet", "no_vest"}:
            color = (220, 80, 60)
        else:
            color = (90, 100, 120)

        x1, y1, x2, y2 = map(int, (box.x1, box.y1, box.x2, box.y2))
        draw.rectangle((x1, y1, x2, y2), outline=color, width=3)
        draw_label(draw, x1, y1, f"{box.label} {box.confidence:.2f}", color)

    return annotated_image


def empty_outputs(message: str) -> tuple[Image.Image | None, str, str, list[list[str]], list[list[str]], str]:
    return None, message, "", [], [], ""


def detect_ppe(
    image: Image.Image,
    person_confidence: float,
    ppe_confidence: float,
    iou_threshold: float,
) -> tuple[Image.Image | None, str, str, list[list[str]], list[list[str]], str]:
    if image is None:
        return empty_outputs("### PPE Check Result\n\nPlease upload an image.")

    try:
        ppe_model = load_ppe_model()
        person_model = load_person_model()
    except Exception as exc:
        return (
            image,
            "### Model Loading Error\n\n"
            "The detection model could not be loaded. Check the Hugging Face Space logs, "
            "network access, or model repository settings.\n\n"
            f"`{type(exc).__name__}: {exc}`",
            "",
            [],
            [],
            "",
        )

    try:
        rgb_image = image.convert("RGB")
        image_array = np.array(rgb_image)

        person_results = person_model.predict(
            source=image_array,
            conf=person_confidence,
            iou=iou_threshold,
            imgsz=IMAGE_SIZE,
            classes=[0],
            verbose=False,
        )
        ppe_results = ppe_model.predict(
            source=image_array,
            conf=ppe_confidence,
            iou=iou_threshold,
            imgsz=IMAGE_SIZE,
            verbose=False,
        )

        workers = extract_person_boxes(person_results[0])
        ppe_boxes = extract_ppe_boxes(ppe_results[0])
        worker_ppe = assign_ppe_to_workers(workers, ppe_boxes)
        ppe_assignment = build_ppe_assignment_lookup(worker_ppe)
        associated_ppe_boxes = sorted(
            {box for assigned_boxes in worker_ppe.values() for box in assigned_boxes},
            key=lambda box: (box.y1, box.x1),
        )
        assessments = assess_workers(workers, worker_ppe)

        annotated_image = draw_detection_boxes(rgb_image, workers, ppe_boxes, assessments, ppe_assignment)
        dashboard = "\n".join(
            [
                build_status_markdown(assessments),
                build_ppe_markdown(assessments, associated_ppe_boxes),
                build_worker_summary_markdown(assessments),
            ]
        )

        return (
            annotated_image,
            dashboard,
            "### Per-Worker PPE Assessment",
            build_worker_table(assessments),
            build_detection_table(workers, ppe_boxes, ppe_assignment),
            f"Processed {len(workers)} worker(s), {len(associated_ppe_boxes)} associated PPE detection(s), and {len(ppe_boxes) - len(associated_ppe_boxes)} unassigned PPE detection(s).",
        )
    except Exception as exc:
        return (
            image,
            "### Detection Error\n\n"
            "The image could not be processed. Try another image or adjust the thresholds.\n\n"
            f"`{type(exc).__name__}: {exc}`",
            "",
            [],
            [],
            "",
        )


def get_example_images() -> list[str]:
    sample_dir = Path("sample_images")
    if not sample_dir.exists():
        return []
    return [
        str(path)
        for path in sample_dir.iterdir()
        if path.suffix.lower() in {".jpg", ".jpeg", ".png", ".webp"}
    ]


CUSTOM_CSS = """
.gradio-container {
    max-width: 1240px !important;
}
.app-subtitle {
    color: #475569;
    font-size: 1.03rem;
    line-height: 1.6;
}
.status-note textarea {
    font-weight: 650;
}
"""

THEME = gr.themes.Soft(primary_hue="blue", neutral_hue="slate")
BLOCKS_KWARGS = {"title": APP_NAME}
LAUNCH_KWARGS = {}

if "theme" in signature(gr.Blocks).parameters:
    BLOCKS_KWARGS.update({"theme": THEME, "css": CUSTOM_CSS})
else:
    LAUNCH_KWARGS.update({"theme": THEME, "css": CUSTOM_CSS})


with gr.Blocks(**BLOCKS_KWARGS) as demo:
    gr.Markdown(
        f"""
# {APP_NAME}

<div class="app-subtitle">
Industrial safety monitoring prototype for factories, warehouses, and construction sites.
Upload a workplace image to detect workers, associate PPE with each worker, and generate
an operational OK / NG safety summary.
</div>
"""
    )

    with gr.Row():
        with gr.Column(scale=5):
            input_image = gr.Image(
                type="pil",
                label="Upload workplace safety image",
                sources=["upload"],
                height=430,
            )
            with gr.Row():
                person_confidence_slider = gr.Slider(
                    minimum=0.1,
                    maximum=0.9,
                    value=PERSON_CONFIDENCE_THRESHOLD,
                    step=0.05,
                    label="Worker confidence threshold",
                )
                ppe_confidence_slider = gr.Slider(
                    minimum=0.1,
                    maximum=0.9,
                    value=PPE_CONFIDENCE_THRESHOLD,
                    step=0.05,
                    label="PPE confidence threshold",
                )
            iou_slider = gr.Slider(
                minimum=0.1,
                maximum=0.9,
                value=IOU_THRESHOLD,
                step=0.05,
                label="IoU threshold",
            )
            detect_button = gr.Button("Run Site Safety Check", variant="primary", size="lg")

        with gr.Column(scale=7):
            output_image = gr.Image(label="Annotated detection result", height=430)
            status_note = gr.Textbox(label="Processing note", lines=1, elem_classes=["status-note"])

    with gr.Tabs():
        with gr.Tab("Safety Dashboard"):
            dashboard_markdown = gr.Markdown()
        with gr.Tab("Worker Assessment"):
            worker_title = gr.Markdown("### Per-Worker PPE Assessment")
            worker_table = gr.Dataframe(
                headers=[
                    "Worker",
                    "Helmet",
                    "Safety Vest",
                    "Overall",
                    "Missing Items",
                    "Worker Conf.",
                    "Helmet Conf.",
                    "Vest Conf.",
                ],
                datatype=["str", "str", "str", "str", "str", "str", "str", "str"],
                interactive=False,
                wrap=True,
            )
        with gr.Tab("Detection Details"):
            detection_table = gr.Dataframe(
                headers=["Type", "Label", "Confidence", "Assigned Worker", "Box"],
                datatype=["str", "str", "str", "str", "str"],
                interactive=False,
                wrap=True,
            )

    examples = get_example_images()
    if examples:
        gr.Examples(examples=examples, inputs=input_image, label="Sample Images")

    detect_button.click(
        fn=detect_ppe,
        inputs=[input_image, person_confidence_slider, ppe_confidence_slider, iou_slider],
        outputs=[output_image, dashboard_markdown, worker_title, worker_table, detection_table, status_note],
    )


if __name__ == "__main__":
    demo.launch(**LAUNCH_KWARGS)