from __future__ import annotations

import re
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import Any


_MODEL_ID_SUFFIX_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,63}$")


@dataclass(frozen=True)
class ServingSpec:
    model_id: str
    model_source: str
    num_replicas: int = 1
    gpus_per_replica: int = 1
    engine_kwargs: dict[str, Any] | None = None


@dataclass(frozen=True)
class ResolvedServingSpec:
    user_id: str
    model_id_suffix: str
    model_id_prefix: str
    model_id: str
    model_source: str
    num_replicas: int
    gpus_per_replica: int
    engine_kwargs: dict[str, Any] | None


def validate_model_id_suffix(suffix: str) -> None:
    if not isinstance(suffix, str):
        raise ValueError("model_id must be a string")
    s = suffix.strip()
    if s != suffix:
        raise ValueError("model_id must not contain leading/trailing whitespace")
    if not s:
        raise ValueError("model_id is required")
    if not _MODEL_ID_SUFFIX_RE.match(s):
        raise ValueError("model_id must match regex: ^[A-Za-z0-9][A-Za-z0-9._-]{0,63}$")
    if ".." in s:
        raise ValueError("model_id must not contain '..'")


def make_model_id_prefix(*, user_id: str, now_utc: datetime | None = None) -> str:
    if not user_id or not isinstance(user_id, str):
        raise ValueError("user_id is required")
    if "/" in user_id:
        raise ValueError("user_id must not contain '/'")

    dt = now_utc or datetime.now(timezone.utc)
    if dt.tzinfo is None:
        dt = dt.replace(tzinfo=timezone.utc)
    dt = dt.astimezone(timezone.utc)
    stamp = dt.strftime("%Y%m%d%H%M")
    return f"{user_id}-{stamp}"


def expand_home_macros(*, user_id: str, text: str) -> str:
    if not isinstance(text, str):
        raise ValueError("model_source must be a string")
    if not text:
        raise ValueError("model_source is required")

    out = text
    out = out.replace("$HOME/common/hf", "/private/hf")
    out = out.replace("$HOME/common/datasets", "/private/datasets")
    out = out.replace("$HOME", f"/private/users/{user_id}")
    return out


def validate_model_source_path(*, user_id: str, model_source: str) -> None:
    if not isinstance(model_source, str):
        raise ValueError("model_source must be a string")
    if not model_source.startswith("/"):
        raise ValueError("model_source must be an absolute path")
    if not model_source.startswith("/private/"):
        raise ValueError("model_source must be under /private")
    if "\x00" in model_source:
        raise ValueError("model_source contains null byte")
    parts = [p for p in model_source.split("/") if p]
    if any(p == ".." for p in parts):
        raise ValueError("model_source must not contain '..'")

    allowed_user_prefix = f"/private/users/{user_id}/"
    allowed = model_source.startswith("/private/hf/") or model_source.startswith(allowed_user_prefix)
    if not allowed:
        raise PermissionError("model_source is not allowed (must be under /private/hf or your /private/users/<user_id>)")


def parse_serving_spec(obj: Any) -> ServingSpec:
    if not isinstance(obj, dict):
        raise ValueError("serving spec must be a mapping")

    model_id = obj.get("model_id")
    model_source = obj.get("model_source")
    num_replicas = obj.get("num_replicas", 1)
    gpus_per_replica = obj.get("gpus_per_replica", 1)
    engine_kwargs = obj.get("engine_kwargs", None)

    if not isinstance(model_id, str):
        raise ValueError("missing required field: model_id")
    validate_model_id_suffix(model_id)

    if not isinstance(model_source, str) or not model_source:
        raise ValueError("missing required field: model_source")

    if not isinstance(num_replicas, int) or num_replicas < 1:
        raise ValueError("num_replicas must be an integer >= 1")
    if not isinstance(gpus_per_replica, int) or gpus_per_replica < 1:
        raise ValueError("gpus_per_replica must be an integer >= 1")

    if engine_kwargs is not None and not isinstance(engine_kwargs, dict):
        raise ValueError("engine_kwargs must be a mapping when provided")

    return ServingSpec(
        model_id=model_id,
        model_source=model_source,
        num_replicas=num_replicas,
        gpus_per_replica=gpus_per_replica,
        engine_kwargs=engine_kwargs,
    )


def resolve_serving_spec(*, spec: ServingSpec, user_id: str, now_utc: datetime | None = None) -> ResolvedServingSpec:
    validate_model_id_suffix(spec.model_id)
    prefix = make_model_id_prefix(user_id=user_id, now_utc=now_utc)
    full_model_id = f"{prefix}-{spec.model_id}"

    resolved_source = expand_home_macros(user_id=user_id, text=spec.model_source)
    validate_model_source_path(user_id=user_id, model_source=resolved_source)

    return ResolvedServingSpec(
        user_id=user_id,
        model_id_suffix=spec.model_id,
        model_id_prefix=prefix,
        model_id=full_model_id,
        model_source=resolved_source,
        num_replicas=spec.num_replicas,
        gpus_per_replica=spec.gpus_per_replica,
        engine_kwargs=spec.engine_kwargs,
    )
