from __future__ import annotations

import hashlib
import time
import uuid
from dataclasses import dataclass
from threading import Lock
from typing import Any, Mapping

import cloudpickle

from flownet.core.flow_program import (
    normalize_flow_program_version,
    normalize_flow_uri,
)


def _normalize_flow_uri(flow_uri: str) -> str:
    return normalize_flow_uri(flow_uri)


def _normalize_version(version: str) -> str:
    return normalize_flow_program_version(version)


def _normalize_mapping(name: str, value: Mapping[str, Any] | None) -> dict[str, Any]:
    if value is None:
        return {}
    if not isinstance(value, Mapping):
        raise ValueError(f"{name} must be a mapping when provided.")
    return dict(value)


def _payload_hash(payload: Any) -> str:
    try:
        blob = cloudpickle.dumps(payload)
        return hashlib.sha256(blob).hexdigest()
    except Exception:
        material = repr(payload).encode("utf-8", errors="ignore")
        return hashlib.sha256(material).hexdigest()


def _tags_from_metadata(metadata: Mapping[str, Any]) -> dict[str, str]:
    raw = metadata.get("tags")
    if not isinstance(raw, Mapping):
        return {}
    normalized: dict[str, str] = {}
    for key, value in raw.items():
        k = str(key).strip()
        if not k:
            continue
        normalized[k] = str(value)
    return normalized


def _match_tags(metadata: Mapping[str, Any], tags: Mapping[str, Any] | None) -> bool:
    if not tags:
        return True
    if not isinstance(tags, Mapping):
        return False
    available = _tags_from_metadata(metadata)
    for key, expected in tags.items():
        if available.get(str(key)) != str(expected):
            return False
    return True


@dataclass(frozen=True)
class FlowProgramRecord:
    program_id: str
    flow_uri: str
    version: str
    program_hash: str
    payload: Any
    metadata: dict[str, Any]
    io_contract: dict[str, Any]
    registered_at: float
    updated_at: float
    status: str = "active"

    def flow_id(self) -> str:
        return f"{self.flow_uri}@{self.version}"

    def to_public_dict(self) -> dict[str, Any]:
        return {
            "program_id": self.program_id,
            "flow_uri": self.flow_uri,
            "version": self.version,
            "flow_id": self.flow_id(),
            "program_hash": self.program_hash,
            "metadata": dict(self.metadata),
            "io_contract": dict(self.io_contract),
            "registered_at": self.registered_at,
            "updated_at": self.updated_at,
            "status": self.status,
        }


class FlowProgramCatalog:
    def __init__(self):
        self._lock = Lock()
        self._records: dict[tuple[str, str], FlowProgramRecord] = {}

    def register(
        self,
        *,
        flow_uri: str,
        version: str,
        payload: Any,
        metadata: Mapping[str, Any] | None = None,
        io_contract: Mapping[str, Any] | None = None,
    ) -> FlowProgramRecord:
        normalized_uri = _normalize_flow_uri(flow_uri)
        normalized_version = _normalize_version(version)
        key = (normalized_uri, normalized_version)
        metadata_dict = _normalize_mapping("metadata", metadata)
        io_contract_dict = _normalize_mapping("io_contract", io_contract)
        now = time.time()
        record = FlowProgramRecord(
            program_id=uuid.uuid4().hex,
            flow_uri=normalized_uri,
            version=normalized_version,
            program_hash=_payload_hash(payload),
            payload=payload,
            metadata=metadata_dict,
            io_contract=io_contract_dict,
            registered_at=now,
            updated_at=now,
            status="active",
        )
        with self._lock:
            if key in self._records:
                raise ValueError(
                    f"FlowProgram already exists: {normalized_uri}@{normalized_version}. "
                    "Use delete_flow_program(...) then register again."
                )
            self._records[key] = record
        return record

    def get(self, flow_uri: str, version: str | None = None) -> FlowProgramRecord | None:
        normalized_uri = _normalize_flow_uri(flow_uri)
        with self._lock:
            if version is not None:
                normalized_version = _normalize_version(version)
                return self._records.get((normalized_uri, normalized_version))
            candidates = [
                record
                for (uri, _version), record in self._records.items()
                if uri == normalized_uri
            ]
        if not candidates:
            return None
        return max(candidates, key=lambda item: item.registered_at)

    def list(
        self,
        *,
        prefix: str | None = None,
        tags: Mapping[str, Any] | None = None,
        limit: int | None = None,
    ) -> list[FlowProgramRecord]:
        with self._lock:
            records = list(self._records.values())
        if prefix:
            records = [record for record in records if record.flow_uri.startswith(prefix)]
        records = [record for record in records if _match_tags(record.metadata, tags)]
        records.sort(key=lambda item: item.registered_at, reverse=True)
        if isinstance(limit, int) and limit >= 0:
            records = records[:limit]
        return records

    def delete(self, flow_uri: str, version: str | None = None) -> int:
        normalized_uri = _normalize_flow_uri(flow_uri)
        removed = 0
        with self._lock:
            if version is not None:
                normalized_version = _normalize_version(version)
                key = (normalized_uri, normalized_version)
                if key in self._records:
                    self._records.pop(key, None)
                    return 1
                return 0
            keys_to_remove = [
                key
                for key in self._records.keys()
                if key[0] == normalized_uri
            ]
            for key in keys_to_remove:
                self._records.pop(key, None)
                removed += 1
        return removed


__all__ = ["FlowProgramRecord", "FlowProgramCatalog"]
