# backend\ai_server.py

import json
import os
import secrets
from pathlib import Path
from typing import List, Optional

from fastapi import Depends, FastAPI, HTTPException, Request, status
from pydantic import BaseModel
from ultralytics import YOLO


BASE_DIR = Path(__file__).resolve().parent


def existing_file(path: Path) -> Optional[Path]:
    try:
        if path.exists() and path.is_file() and path.stat().st_size > 0:
            return path
    except OSError:
        pass
    return None


def resolve_training_root() -> Path:
    env_root = os.environ.get("TRAINING_ROOT", "").strip()
    if env_root:
        root = Path(env_root).expanduser().resolve()
        root.mkdir(parents=True, exist_ok=True)
        return root

    candidates = [
        # Wenn ai_server.py aus backend/ läuft:
        BASE_DIR / "generated" / "training",

        # Wenn ai_server.py aus backend/ml/ laufen würde:
        BASE_DIR.parent / "generated" / "training",

        # Wenn ai_server.py embedded aus Temp läuft, aber backendRoot als cwd gesetzt wurde:
        Path.cwd() / "generated" / "training",

        # Wenn Working Directory Projektroot ist:
        Path.cwd() / "backend" / "generated" / "training",
    ]

    for root in candidates:
        if (
            existing_file(root / "detection_labels.json")
            or existing_file(root / "detector" / "model" / "best.pt")
        ):
            root.mkdir(parents=True, exist_ok=True)
            return root.resolve()

    # Fallback: Server soll trotzdem starten.
    root = (Path.cwd() / "generated" / "training").resolve()
    root.mkdir(parents=True, exist_ok=True)
    return root


TRAINING_ROOT = resolve_training_root()
DEFAULT_MODEL_PATH = TRAINING_ROOT / "detector" / "model" / "best.pt"


def resolve_detection_labels_path() -> Path:
    env_path = os.environ.get("DETECTION_LABELS_PATH", "").strip()
    if env_path:
        p = Path(env_path).expanduser().resolve()
        if existing_file(p):
            return p
        raise RuntimeError(f"DETECTION_LABELS_PATH not found: {p}")

    p = TRAINING_ROOT / "detection_labels.json"
    if existing_file(p):
        return p.resolve()

    raise RuntimeError(f"detection_labels.json not found: {p}")


def resolve_model_path() -> str:
    env_path = os.environ.get("YOLO_MODEL", "").strip()
    if env_path:
        p = Path(env_path).expanduser().resolve()
        if existing_file(p):
            return str(p)
        raise RuntimeError(f"YOLO_MODEL not found: {p}")

    if existing_file(DEFAULT_MODEL_PATH):
        return str(DEFAULT_MODEL_PATH)

    raise RuntimeError(f"YOLO model not found: {DEFAULT_MODEL_PATH}")


# Server darf auch ohne Labels/Model starten.
DETECTION_LABELS_PATH: Optional[Path] = None

LABEL_GROUPS = {
    "people": set(),
    "sexPositions": {"unknown"},
    "bodyParts": set(),
    "objects": set(),
    "clothing": set(),
}

BODY_LABELS = LABEL_GROUPS["bodyParts"]
OBJECT_LABELS = LABEL_GROUPS["objects"]
CLOTHING_LABELS = LABEL_GROUPS["clothing"]
POSITION_LABELS = set()

PERSON_LABELS = {
    "person_male",
    "person_female",
}

_MODEL_PATH = ""
_MODEL_ERROR = ""
_LABEL_ERROR = ""

_DEVICE = os.environ.get("YOLO_DEVICE", "")
_CONF = float(os.environ.get("YOLO_CONF", "0.25"))
_BATCH = int(os.environ.get("YOLO_BATCH", "16"))
_IMGSZ = int(os.environ.get("YOLO_IMGSZ", "640"))
_HALF = os.environ.get("YOLO_HALF", "0").lower() in {"1", "true", "yes", "on"}

model = None

app = FastAPI()

AI_SERVER_TOKEN = os.environ.get("AI_SERVER_TOKEN", "").strip()
AI_SERVER_AUTH_REQUIRED = os.environ.get("AI_SERVER_AUTH_REQUIRED", "1").strip().lower() not in {
    "0",
    "false",
    "no",
    "off",
}


def require_ai_server_auth(request: Request) -> None:
    if not AI_SERVER_AUTH_REQUIRED:
        return

    if not AI_SERVER_TOKEN:
        raise HTTPException(
            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
            detail="AI server auth token not configured",
        )

    auth = request.headers.get("authorization", "").strip()
    header_token = request.headers.get("x-ai-server-token", "").strip()

    provided = ""

    if auth.lower().startswith("bearer "):
        provided = auth[7:].strip()
    elif header_token:
        provided = header_token

    if not provided or not secrets.compare_digest(provided, AI_SERVER_TOKEN):
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="unauthorized",
            headers={"WWW-Authenticate": "Bearer"},
        )


class PredictBatchRequest(BaseModel):
    paths: List[str]
    detectorOnly: bool = False
    imageSize: int = 640
    model: Optional[str] = None


def empty_prediction(source: str = "model_missing") -> dict:
    return {
        "modelAvailable": False,
        "source": source,
        "sexPosition": "unknown",
        "sexPositionScore": 0.0,
        "peoplePresent": [],
        "bodyPartsPresent": [],
        "objectsPresent": [],
        "clothingPresent": [],
        "boxes": [],
    }


def load_label_groups_safe() -> None:
    global DETECTION_LABELS_PATH
    global LABEL_GROUPS
    global BODY_LABELS
    global OBJECT_LABELS
    global CLOTHING_LABELS
    global POSITION_LABELS
    global PERSON_LABELS
    global _LABEL_ERROR

    try:
        path = resolve_detection_labels_path()
        DETECTION_LABELS_PATH = path

        with path.open("r", encoding="utf-8") as f:
            data = json.load(f)

        LABEL_GROUPS = {
            "people": set(
                str(x).strip().lower()
                for x in data.get("people", [])
                if str(x).strip()
            ),
            "sexPositions": set(
                str(x).strip().lower()
                for x in data.get("sexPositions", [])
                if str(x).strip()
            ),
            "bodyParts": set(
                str(x).strip().lower()
                for x in data.get("bodyParts", [])
                if str(x).strip()
            ),
            "objects": set(
                str(x).strip().lower()
                for x in data.get("objects", [])
                if str(x).strip()
            ),
            "clothing": set(
                str(x).strip().lower()
                for x in data.get("clothing", [])
                if str(x).strip()
            ),
        }

        if not LABEL_GROUPS["sexPositions"]:
            LABEL_GROUPS["sexPositions"] = {"unknown"}

        _LABEL_ERROR = ""

    except Exception as exc:
        DETECTION_LABELS_PATH = None
        _LABEL_ERROR = str(exc)

        LABEL_GROUPS = {
            "people": set(),
            "sexPositions": {"unknown"},
            "bodyParts": set(),
            "objects": set(),
            "clothing": set(),
        }

    BODY_LABELS = LABEL_GROUPS["bodyParts"]
    OBJECT_LABELS = LABEL_GROUPS["objects"]
    CLOTHING_LABELS = LABEL_GROUPS["clothing"]

    POSITION_LABELS = {
        label for label in LABEL_GROUPS["sexPositions"]
        if label and label != "unknown"
    }

    PERSON_LABELS = {
        label for label in LABEL_GROUPS["people"]
        if label
    }


def get_model():
    global model
    global _MODEL_PATH
    global _MODEL_ERROR

    if model is not None:
        return model

    try:
        path = resolve_model_path()
        loaded = YOLO(path)

        model = loaded
        _MODEL_PATH = path
        _MODEL_ERROR = ""

        # Labels erst laden, wenn Inference wirklich gebraucht wird.
        load_label_groups_safe()

        return model

    except Exception as exc:
        model = None
        _MODEL_PATH = ""
        _MODEL_ERROR = str(exc)
        return None


def scored(label: str, score: float) -> dict:
    return {
        "label": label,
        "score": float(score),
    }


def best_score(items: list[dict], label: str, score: float) -> None:
    for item in items:
        if item["label"] == label:
            if score > item["score"]:
                item["score"] = float(score)
            return

    items.append(scored(label, score))


def prediction_from_result(result) -> dict:
    names = result.names or {}

    boxes_out = []
    people_present = []
    body_parts = []
    objects = []
    clothing = []

    sex_position = "unknown"
    sex_position_score = 0.0

    if result.boxes is not None:
        xywhn = result.boxes.xywhn.cpu().tolist()
        cls_values = result.boxes.cls.cpu().tolist()
        conf_values = result.boxes.conf.cpu().tolist()

        for box_xywhn, cls_id, conf in zip(xywhn, cls_values, conf_values):
            label = str(names.get(int(cls_id), int(cls_id))).strip().lower()
            score = float(conf)

            if not label:
                continue

            cx, cy, w, h = [float(v) for v in box_xywhn]

            x = max(0.0, min(1.0, cx - w / 2.0))
            y = max(0.0, min(1.0, cy - h / 2.0))
            w = max(0.0, min(1.0 - x, w))
            h = max(0.0, min(1.0 - y, h))

            is_person = label in PERSON_LABELS
            is_body = label in BODY_LABELS
            is_object = label in OBJECT_LABELS
            is_clothing = label in CLOTHING_LABELS
            is_position = label in POSITION_LABELS

            if is_position:
                if score > sex_position_score:
                    sex_position = label
                    sex_position_score = score
                continue

            if not (is_person or is_body or is_object or is_clothing):
                continue

            boxes_out.append({
                "label": label,
                "score": score,
                "x": x,
                "y": y,
                "w": w,
                "h": h,
            })

            if is_person:
                best_score(people_present, label, score)

            if is_body:
                best_score(body_parts, label, score)

            if is_object:
                best_score(objects, label, score)

            if is_clothing:
                best_score(clothing, label, score)

    return {
        "modelAvailable": True,
        "source": f"yolo-server:{Path(_MODEL_PATH).name}",
        "sexPosition": sex_position,
        "sexPositionScore": sex_position_score,
        "peoplePresent": people_present,
        "bodyPartsPresent": body_parts,
        "objectsPresent": objects,
        "clothingPresent": clothing,
        "boxes": boxes_out,
    }


@app.post("/predict-batch", dependencies=[Depends(require_ai_server_auth)])
def predict_batch(req: PredictBatchRequest):
    paths = [str(path).strip() for path in req.paths if str(path).strip()]
    if not paths:
        return {
            "ok": False,
            "predictions": [],
            "error": "no paths supplied",
        }

    current_model = get_model()
    if current_model is None:
        return {
            "ok": True,
            "predictions": [empty_prediction("model_missing") for _ in paths],
            "error": _MODEL_ERROR or f"YOLO model not found: {DEFAULT_MODEL_PATH}",
        }

    if DETECTION_LABELS_PATH is None or _LABEL_ERROR:
        return {
            "ok": False,
            "predictions": [],
            "error": f"detection labels missing: {_LABEL_ERROR}",
        }

    imgsz = int(req.imageSize or _IMGSZ or 640)

    try:
        results = current_model.predict(
            source=paths,
            imgsz=imgsz,
            conf=_CONF,
            batch=_BATCH,
            device=_DEVICE or None,
            half=_HALF,
            verbose=False,
        )

        predictions = [prediction_from_result(result) for result in results]

        return {
            "ok": True,
            "predictions": predictions,
        }

    except Exception as exc:
        return {
            "ok": False,
            "predictions": [],
            "error": str(exc),
        }


@app.get("/health", dependencies=[Depends(require_ai_server_auth)])
def health():
    current_model = get_model()
    names = getattr(current_model, "names", {}) or {} if current_model is not None else {}

    return {
        "ok": True,
        "ready": current_model is not None,
        "modelAvailable": current_model is not None,
        "model": _MODEL_PATH,
        "modelError": _MODEL_ERROR,
        "expectedModel": str(DEFAULT_MODEL_PATH),
        "trainingRoot": str(TRAINING_ROOT),
        "classCount": len(names),
        "classes": list(names.values())[:80] if isinstance(names, dict) else names,
        "labelConfig": str(DETECTION_LABELS_PATH) if DETECTION_LABELS_PATH else "",
        "labelError": _LABEL_ERROR,
    }

@app.post("/reload", dependencies=[Depends(require_ai_server_auth)])
def reload_model():
    global model
    global _MODEL_PATH
    global _MODEL_ERROR
    global DETECTION_LABELS_PATH

    model = None
    _MODEL_PATH = ""
    _MODEL_ERROR = ""
    DETECTION_LABELS_PATH = None

    current_model = get_model()
    names = getattr(current_model, "names", {}) or {} if current_model is not None else {}

    return {
        "ok": current_model is not None,
        "ready": current_model is not None,
        "modelAvailable": current_model is not None,
        "model": _MODEL_PATH,
        "modelError": _MODEL_ERROR,
        "expectedModel": str(DEFAULT_MODEL_PATH),
        "trainingRoot": str(TRAINING_ROOT),
        "classCount": len(names),
        "classes": list(names.values())[:80] if isinstance(names, dict) else names,
        "labelConfig": str(DETECTION_LABELS_PATH) if DETECTION_LABELS_PATH else "",
        "labelError": _LABEL_ERROR,
    }