"""
Manifest data structures and signing utilities.

This module defines the canonical schema for a Score Vision Manifest.
A Manifest is a cryptographically signed, content-addressed rulebook
for a single evaluation window. It specifies Elements, multi-pillar metrics,
performance baselines (baseline_theta), difficulty weights (beta),
minimum performance margins (delta_floor), latency gates, service rates,
and TEE-related trust parameters.

- baseline_theta (θₑ): Minimum Element-specific performance required to earn rewards.
- delta_floor: Minimum margin above baseline before applying difficulty weighting.
- beta (βₑ): Difficulty multiplier scaling rewards based on Element complexity.
"""

import os
from enum import Enum
from functools import cached_property
from hashlib import sha256
from json import dumps, loads
from pathlib import Path
import aiohttp
import hashlib

from cv2 import imread
from nacl.exceptions import BadSignatureError
from nacl.signing import SigningKey, VerifyKey
from numpy import ndarray
from pydantic import BaseModel, Field, model_validator
from ruamel.yaml import YAML
from urllib.parse import urlparse, urljoin

yaml = YAML()
yaml.default_flow_style = False

# ------------------------------------------------------------
# ENUMS
# ------------------------------------------------------------


class NormType(str, Enum):
    """Preprocessing normalisation modes."""

    RGB_01 = "rgb-01"
    RGB_255 = "rgb-255"
    NONE = "none"


class PillarName(str, Enum):
    """
    Multi-pillar metrics used by Elements.
    (See Appendix E for naming conventions).
    """

    IOU = "iou"
    COUNT = "count"
    PALETTE = "palette"
    SMOOTHNESS = "smoothness"
    ROLE = "role"
    KEYPOINTS_IOU = "keypoints_iou" 
    MAP50 = "map50"
    PRECISION = "precision"
    RECALL = "recall"
    SOCCER_ACTION = "soccer_action"
    CRICKET_SCORING = "cricket_scoring"
    TCG_GRADING = "tcg_grading"
    FALSE_POSITIVE = "false_positive"


class ElementPrefix(str, Enum):
    """Prefix to Element Names
    e.g.
    - PlayerDetect_v1: Object detection and tracking at. Outputs include bounding boxes, tracking IDs, team assignments, and role classifications.
    - BallDetect_v1: Small object tracking. Handles fast-moving objects with motion blur and frequent occlusions.
    - PitchCalib_v1: Geometric calibration. Outputs keypoint locations and homography matrices for image-to-field coordinate transformation.
    """

    PITCH_CALIBRATION = "PitchCalib"
    OBJECT_DETECTION = "Detect"
    # NOTE:Player Detection is a special case of Object Detection that adds team classification
    PLAYER_DETECTION = "PlayerDetect"
    OBJECT_TRACKING = "Track"
    SCENE_CLASSIFICATION = "ClassifyImage"
    OBJECT_CLASSIFICATION = "ClassifyObject"


class ChallengeType(str, Enum):
    FOOTBALL = "football"
    CRICKET = "cricket"
    BASKETBALL = "basketball"


class GroundTruthType(str, Enum):
    SOCCER_ACTION = "soccer_action"
    CRICKET_DELIVERY = "cricket_delivery"
    TCG_GRADING = "tcg_grading"


# ------------------------------------------------------------
# DATA CLASSES
# ------------------------------------------------------------


class KeypointTemplate(BaseModel):
    filename: str = Field(..., description="name of keypoints template image file")
    keypoints_on_template: list[tuple[int, int]] = Field(
        ...,
        description="the pixel coordinates of each keypoint on the template image file",
    )
    keypoint_index_bottom_left: int = Field(..., ge=0)
    keypoint_index_bottom_right: int = Field(..., ge=0)
    keypoint_index_top_left: int = Field(..., ge=0)
    keypoint_index_top_right: int = Field(..., ge=0)

    @property
    def bottom_left(self) -> tuple[int, int]:
        return self.keypoints_on_template[self.keypoint_index_bottom_left]

    @property
    def bottom_right(self) -> tuple[int, int]:
        return self.keypoints_on_template[self.keypoint_index_bottom_right]

    @property
    def top_left(self) -> tuple[int, int]:
        return self.keypoints_on_template[self.keypoint_index_top_left]

    @property
    def top_right(self) -> tuple[int, int]:
        return self.keypoints_on_template[self.keypoint_index_top_right]

    @property
    def n_keypoints(self) -> int:
        return len(self.keypoints_on_template)

    @property
    def template(self) -> ndarray:
        current_dir = Path(__file__).parent.parent
        path_template = str(current_dir / "keypoints_templates" / self.filename)
        image = imread(path_template)
        if image is None:
            raise ValueError(f"No image file found for template at: {path_template}")
        return image

    @model_validator(mode="after")
    def validate_template_and_indices(self):
        self.template
        try:
            self.bottom_left
            self.bottom_right
            self.top_left
            self.top_right
        except IndexError as e:
            raise ValueError(f"index must be within 0 and {self.n_keypoints - 1}: {e}")
        return self


KEYPOINT_TEMPLATES: dict[ChallengeType, KeypointTemplate] = {
    ChallengeType.CRICKET: KeypointTemplate(
        filename="cricket_pitch_template.png",
        keypoints_on_template=[
            (35, 1),  # 1
            (325, 1),  # 2
            # -------------
            (35, 120),  # 3
            (325, 120),  # 4
            # -------------
            (35, 235),  # 5
            (325, 235),  # 6
            # ===============
            (35, 1160),  # 7
            (325, 1160),  # 8
            # -------------
            (35, 1275),  # 9
            (325, 1275),  # 10
            # -------------
            (35, 1392),  # 11
            (325, 1392),  # 12
            # ===============
            (175, 120),  # 13
            (175, 1275),  # 14
        ],
        keypoint_index_bottom_left=10,
        keypoint_index_bottom_right=11,
        keypoint_index_top_left=0,
        keypoint_index_top_right=1,
    ),
    ChallengeType.FOOTBALL: KeypointTemplate(
        filename="football_pitch_template.png",
        keypoints_on_template=[
            (5, 5),  # 1
            (5, 140),  # 2
            (5, 250),  # 3
            (5, 430),  # 4
            (5, 540),  # 5
            (5, 675),  # 6
            # -------------
            (55, 250),  # 7
            (55, 430),  # 8
            # -------------
            (110, 340),  # 9
            # -------------
            (165, 140),  # 10
            (165, 270),  # 11
            (165, 410),  # 12
            (165, 540),  # 13
            # -------------
            (527, 5),  # 14
            (527, 253),  # 15
            (527, 433),  # 16
            (527, 675),  # 17
            # -------------
            (888, 140),  # 18
            (888, 270),  # 19
            (888, 410),  # 20
            (888, 540),  # 21
            # -------------
            (940, 340),  # 22
            # -------------
            (998, 250),  # 23
            (998, 430),  # 24
            # -------------
            (1045, 5),  # 25
            (1045, 140),  # 26
            (1045, 250),  # 27
            (1045, 430),  # 28
            (1045, 540),  # 29
            (1045, 675),  # 30
            # -------------
            (435, 340),  # 31
            (615, 340),  # 32
        ],
        keypoint_index_bottom_left=5,
        keypoint_index_bottom_right=29,
        keypoint_index_top_left=0,
        keypoint_index_top_right=24,
    ),
}


class Preproc(BaseModel):
    """
    Preprocessing parameters applied before evaluation.
    """

    fps: int
    resize_long: int
    norm: NormType


class Metrics(BaseModel):
    """
    Metrics configuration used to score Element performance.
    """

    pillars: dict[PillarName, float]

    @model_validator(mode="after")
    def validate_pillar_weights(self):
        """Ensure the Weights for all pillars sum to 1.0"""
        if sum(self.pillars.values()) != 1.0:
            raise ValueError(f"Weights must sum to 1.0: {self.pillars.values()}")
        return self


class Salt(BaseModel):
    """
    VRF-derived challenge salting parameters ensuring per-validator
    unpredictability. These are always present (default empty lists).
    """

    offsets: list[int] = Field(default_factory=list)
    strides: list[int] = Field(default_factory=list)


class Clip(BaseModel):
    """
    Clip definition used in YAML configs.
    Structured as:
      - hash: "sha256:..."
        weight: 1.0
    """

    hash: str
    weight: float


class Element(BaseModel):
    """
    Atomic capability definition (e.g., PlayerDetect, BallDetect).

    Fields:
    id:                Unique versioned name (e.g., "PlayerDetect_v1@1.0")
    clips:             Optional list of Clip objects
    preproc:           Preprocessing config (fps, resize, normalization)
    metrics:           Metric pillar weights (sum to 1.0)
    latency_p95_ms:    Hard p95 latency gate in milliseconds
    service_rate_fps:  Target real-time service rate (fps)
    pgt_recipe_hash:   Optional immutable hash of the Pseudo-Ground Truth recipe
    baseline_theta:    θₑ — Minimum Element-specific performance required to earn rewards
    delta_floor:       Minimum margin above baseline (Qₑ floor) before applying difficulty weighting
    beta:              βₑ — Difficulty multiplier to scale rewards for harder Elements
    salt:              Per-validator VRF salting offsets and strides for unpredictability
    """

    id: str
    track: str | None = None
    first_block: int | None = None
    window_block: int | None = None
    eval_window: int | float | None = None
    weight: float | None = None
    clips: list[Clip] = Field(default_factory=list)
    metrics: Metrics | None = None
    preproc: Preproc | None = None
    latency_p95_ms: int | None = None
    service_rate_fps: int | None = None
    pgt_recipe_hash: str | None = None
    ground_truth: bool = Field(
        default=False,
        description="If true, use real ground truth fetched from API instead of SAM3 pseudo-GT",
    )
    baseline_theta: float = Field(
        default=0.0, ge=0.0, description="Score threshold for emissions (min 0.0)"
    )
    delta_floor: float = Field(
        default=0.0,
        ge=0.0,
        description=(
            "Minimum margin above baseline (Q_e floor). "
            "Rewards are computed only for improvements above baseline_theta. "
            "This ensures a non-negative minimum contribution before applying beta."
        ),
    )
    beta: float = Field(
        default=1.0,
        ge=0.0,
        description="Difficulty weight for emission allocation (must be positive)",
    )
    max_model_size_mb: float | None = Field(
        default=None,
        ge=0.0,
        description=(
            "Maximum allowed total Hugging Face model repository size in MB for miners "
            "serving this element. Miners above this limit are excluded from registry."
        ),
    )
    onnx_model: bool | None = Field(
        default=None,
        description=(
            "If true, miners serving this element must expose only ONNX model artifacts "
            "in their Hugging Face repository revision."
        ),
    )
    salt: Salt = Field(default_factory=Salt)
    keypoint_template: ChallengeType | None = None
    objects: list[str] | None = None
    challenge_type_version: str | None = None
    groundtruth_type: GroundTruthType | None = None

    def apply_baseline_gate(self, score: float) -> float:
        """Clamp score to positive margin above baseline."""
        return max(score - self.baseline_theta, 0.0)

    def improvement(self, score: float) -> float:
        """Compute improvement with delta_floor applied."""
        return max(self.apply_baseline_gate(score=score), self.delta_floor)

    def apply_difficulty_weight(self, improvement: float) -> float:
        """Scale improvement by beta (difficulty weight)."""
        return self.beta * improvement

    def weight_score(self, score: float) -> float:
        """Convenience function: baseline → delta → beta."""
        return self.apply_difficulty_weight(improvement=self.improvement(score=score))

    @property
    def keypoints(self) -> KeypointTemplate | None:
        return KEYPOINT_TEMPLATES.get(self.keypoint_template)

    @property
    def evaluation_id(self) -> str:
        """
        ID variant used by evaluation logic to infer element family.
        Keeps the full public/API id intact while allowing namespaced ids
        like "namespace/Detect-Person".
        """
        raw_id = str(self.id or "").strip()
        if "/" in raw_id:
            suffix = raw_id.rsplit("/", 1)[-1].strip()
            if suffix:
                return suffix
        return raw_id

    @property
    def category(self) -> ElementPrefix:
        evaluated = self.evaluation_id
        for element_prefix in ElementPrefix:
            if evaluated.startswith(element_prefix.value):
                return element_prefix
        raise ValueError(
            f"Unrecognised element id for category inference: id='{self.id}', evaluation_id='{evaluated}'"
        )

    @model_validator(mode="after")
    def validate_id(self):
        if self.track != "private":
            self.category
        return self


class Tee(BaseModel):
    """
    Trusted Execution Environment parameters.

    Fields:
        trusted_share_gamma: Fraction (γ_trusted ∈ [0,1]) of total emissions
                          allocated to the privacy-preserving Trusted Track.
    """

    trusted_share_gamma: float


class Manifest(BaseModel):
    """
    Canonical Manifest representing one evaluation window.

    Fields:
      window_id:     Unique ID for the window (YYYY-MM-DD)
      elements:      List of Elements with full scoring/configuration
      tee:           TEE trust parameters
      version:       Manifest schema version (e.g., "1.3")
      expiry_block:  Block height after which the manifest expires
      signature:     Base64-encoded Ed25519 signature (optional)

    Methods:
      to_canonical_json() → Stable JSON for hashing/signing.
      sign(private_key)   → Apply Ed25519 signature.
      verify(public_key)  → Verify Ed25519 signature.
      hash                → SHA-256 content-address.
    """

    window_id: str
    version: float
    expiry_block: int
    elements: list[Element]
    tee: Tee
    signature: str | None = None

    @classmethod
    def load_yaml(cls, path: Path) -> "Manifest":
        data = yaml.load(path.read_text())
        return Manifest(**data)

    def get_element(self, id: str) -> Element | None:
        for element in self.elements:
            if element.id == id:
                return element

    def to_canonical_json(self) -> str:
        """
        Produce deterministic canonical JSON suitable for hashing and signing.
        The canonical JSON representation is stable
        and signature-independent
        (i.e. the signature is excluded from hashing and signing)
        so that content hashing is deterministic.
        - Sort Elements lexicographically by ID
        - Exclude signature
        - Compact separators
        - Sorted keys
        """
        payload = self.model_dump(mode="json")
        elements_sorted = sorted(self.elements, key=lambda e: e.id)
        payload["elements"] = [e.model_dump(mode="json") for e in elements_sorted]
        payload.pop("signature", None)
        return dumps(payload, sort_keys=True, separators=(",", ":"), ensure_ascii=False)

    def sign(self, signing_key: SigningKey) -> None:
        """
        Sign the canonical manifest using an Ed25519 private key.
        """
        self.signature = signing_key.sign(
            self.to_canonical_json().encode("utf-8")
        ).signature.hex()

    def verify(self, verify_key: VerifyKey) -> bool:
        """
        Verify the manifest's signature with the corresponding public key.
        """
        if not self.signature:
            raise ValueError("Manifest has no signature to verify.")
        try:
            verify_key.verify(
                self.to_canonical_json().encode("utf-8"), bytes.fromhex(self.signature)
            )
            return True
        except BadSignatureError:
            return False
        except Exception as e:
            raise ValueError(f"Signature verification failed: {e}")

    @cached_property
    def hash(self) -> str:
        """
        Deterministic hashing by
        computing a stable SHA-256 hash of the manifest content
        (This makes the Manifest content-addressable)
        """
        return sha256(self.to_canonical_json().encode("utf-8")).hexdigest()

    @classmethod
    def from_dict(cls, data: dict) -> "Manifest":
        """Rebuild a Manifest (and nested dataclasses) from a plain dict (e.g. JSON)."""
        elements: list[Element] = []
        for e in data.get("elements", []):
            preproc = None
            if e.get("preproc"):
                preproc = Preproc(**e["preproc"])

            metrics = None
            if e.get("metrics"):
                pillars = None
                m = e["metrics"]
                if m.get("pillars"):
                    pillars = Pillars(**m["pillars"])
                metrics = Metrics(pillars=pillars)

            salt = None
            if e.get("salt"):
                salt = Salt(**e["salt"])

            elements.append(
                Element(
                    id=e["id"],
                    track=e.get("track"),
                    first_block=e.get("first_block"),
                    window_block=e.get("window_block"),
                    eval_window=e.get("eval_window"),
                    weight=e.get("weight"),
                    clips=e.get("clips", []),
                    weights=e.get("weights", []),
                    preproc=preproc,
                    metrics=metrics,
                    latency_p95_ms=e.get("latency_p95_ms"),
                    service_rate_fps=e.get("service_rate_fps"),
                    salt=salt,
                    pgt_recipe_hash=e.get("pgt_recipe_hash"),
                    baseline_theta=e.get("baseline_theta"),
                    delta_floor=e.get("delta_floor"),
                    beta=e.get("beta"),
                    groundtruth_type=e.get("groundtruth_type"),
                )
            )

        tee = None
        if data.get("tee"):
            tee = Tee(**data["tee"])

        return cls(
            window_id=data["window_id"],
            elements=elements,
            tee=tee,
            version=data.get("version"),
            expiry_block=data.get("expiry_block"),
            signature=data.get("signature"),
        )

    def save_yaml(self, path: Path) -> None:
        raw = loads(self.to_canonical_json())
        if self.signature:
            raw["signature"] = self.signature
        yaml.dump(raw, path.open("w"))


def get_current_manifest(block_number: int | None = None) -> Manifest:
    """
    Return the current active Manifest.

    Phase 2 version:
    - reads JSON path from env SCOREVISION_MANIFEST_PATH or SV_MANIFEST_PATH
    - optionally checks expiry_block against block_number if provided
    """
    path = os.getenv("SCOREVISION_MANIFEST_PATH") or os.getenv("SV_MANIFEST_PATH")
    if not path:
        raise RuntimeError(
            "No manifest path configured. Set SCOREVISION_MANIFEST_PATH or SV_MANIFEST_PATH."
        )

    manifest = Manifest.load_yaml(path=Path(path))

    if (
        block_number is not None
        and manifest.expiry_block is not None
        and block_number > manifest.expiry_block
    ):
        raise RuntimeError(
            f"Manifest expired at block {manifest.expiry_block}, current block={block_number}."
        )

    return manifest

def _bucket_base(index_url: str) -> str:
    u = urlparse(index_url)
    return f"{u.scheme}://{u.netloc}/"

def _join_key_to_base(index_url: str, key_or_url: str) -> str:
    key_or_url = str(key_or_url).strip()
    while key_or_url.startswith("./"):
        key_or_url = key_or_url[2:]

    if key_or_url.startswith("http://") or key_or_url.startswith("https://"):
        return key_or_url

    if key_or_url.startswith("manako/"):
        return _bucket_base(index_url) + key_or_url

    if key_or_url.startswith("/"):
        return _bucket_base(index_url) + key_or_url.lstrip("/")

    if key_or_url.startswith("manifest/"):
        return _bucket_base(index_url) + key_or_url

    base = index_url.rsplit("/", 1)[0] + "/"
    if base.endswith("/manifest/") and key_or_url.startswith("manifest/"):
        return _bucket_base(index_url) + key_or_url
    return urljoin(base, key_or_url)

async def _http_get_json(url: str, timeout_s: int = 20):
    timeout = aiohttp.ClientTimeout(total=timeout_s)
    async with aiohttp.ClientSession(timeout=timeout) as s:
        async with s.get(url) as r:
            if r.status != 200:
                raise RuntimeError(f"GET {url} -> {r.status}")
            return await r.json()

async def _http_get_text(url: str, timeout_s: int = 30) -> str:
    timeout = aiohttp.ClientTimeout(total=timeout_s)
    async with aiohttp.ClientSession(timeout=timeout) as s:
        async with s.get(url) as r:
            if r.status != 200:
                raise RuntimeError(f"GET {url} -> {r.status}")
            return await r.text()

async def _http_head_meta(url: str, timeout_s: int = 10) -> tuple[str | None, str | None]:
    timeout = aiohttp.ClientTimeout(total=timeout_s)
    async with aiohttp.ClientSession(timeout=timeout) as s:
        async with s.head(url) as r:
            if r.status >= 400:
                return None, None
            return (r.headers.get("ETag"), r.headers.get("Last-Modified"))

def _cache_path_for_url(cache_dir: Path, url: str, suffix: str) -> Path:
    h = hashlib.sha256(url.encode()).hexdigest()[:16]
    name = Path(urlparse(url).path).name or "manifest"
    return cache_dir / f"{name}.{h}.{suffix}"

def _extract_manifest_urls_from_index(index_url: str, idx) -> list[str]:
    keys: list[str] = []
    if isinstance(idx, list):
        keys = [k for k in idx if isinstance(k, str)]
    elif isinstance(idx, dict) and isinstance(idx.get("entries"), list):
        for e in idx["entries"]:
            p = e.get("path")
            if isinstance(p, str):
                keys.append(p)
    return [_join_key_to_base(index_url, k) for k in keys]

def _pick_manifest_url_max_block(urls: list[str]) -> tuple[int, str] | None:
    pairs: list[tuple[int, str]] = []
    for u in urls:
        name = Path(urlparse(u).path).name
        try:
            b = int(name.split("-", 1)[0])
        except Exception:
            continue
        pairs.append((b, u))
    if not pairs:
        return None
    pairs.sort(key=lambda x: x[0])
    return pairs[-1]


def _pick_manifest_url_for_block(
    urls: list[str], block_number: int | None
) -> tuple[int, str] | None:
    pairs: list[tuple[int, str]] = []
    for u in urls:
        name = Path(urlparse(u).path).name
        try:
            b = int(name.split("-", 1)[0])
        except Exception:
            continue
        pairs.append((b, u))

    if not pairs:
        return None

    pairs.sort(key=lambda x: x[0])
    if block_number is None:
        return pairs[-1]

    eligible = [p for p in pairs if p[0] <= block_number]
    if eligible:
        return eligible[-1]

    return pairs[-1]


async def load_manifest_from_public_index(
    index_url: str,
    *,
    block_number: int | None = None,
    cache_dir: Path | None = None,
) -> Manifest:
    idx = await _http_get_json(index_url)
    urls = _extract_manifest_urls_from_index(index_url, idx)
    picked = _pick_manifest_url_for_block(urls, block_number)
    if not picked:
        raise RuntimeError(f"No manifest entries found in index: {index_url}")

    picked_block, manifest_url = picked

    yaml_text: str
    if cache_dir is not None:
        cache_dir.mkdir(parents=True, exist_ok=True)
        out = _cache_path_for_url(cache_dir, manifest_url, "yaml")
        mod = out.with_suffix(".modified")

        etag, lm = await _http_head_meta(manifest_url)
        tag = (etag or lm or "").strip()

        if tag and out.exists() and mod.exists() and mod.read_text().strip() == tag:
            yaml_text = out.read_text()
        else:
            yaml_text = await _http_get_text(manifest_url)
            tmp = out.with_suffix(".tmp")
            tmp.write_text(yaml_text)
            os.replace(tmp, out)
            if tag:
                mod.write_text(tag)
    else:
        yaml_text = await _http_get_text(manifest_url)

    data = yaml.load(yaml_text)
    manifest = Manifest(**data)

    if (
        block_number is not None
        and manifest.expiry_block is not None
        and block_number > manifest.expiry_block
    ):
        raise RuntimeError(
            f"Manifest expired at block {manifest.expiry_block}, current block={block_number}."
        )

    return manifest