"""Top-level SB-DSL convenience copy for imports during tests.

This mirrors the implementation in src/ so tests and consumers can `import spacesafeml_certification_benchmark_and_`.
The implementation is intentionally small and duplicated to avoid packaging complexities in the MVP seed.
"""
from dataclasses import dataclass, asdict, field
from typing import List, Dict, Optional, Any
import json
import hashlib


@dataclass
class AgentCapabilities:
    name: str
    version: str
    compute_arch: Optional[str] = None
    sensors: List[str] = field(default_factory=list)
    actuators: List[str] = field(default_factory=list)


@dataclass
class LocalCapabilities:
    name: str
    capabilities: List[str] = field(default_factory=list)
    metadata: Dict[str, Any] = field(default_factory=dict)


@dataclass
class SafetyContract:
    contract_id: Optional[str] = None
    local_capabilities: List[LocalCapabilities] = field(default_factory=list)
    budgets: Optional["ResourceBudgets"] = None
    data_policy: Optional["DataSharingPolicy"] = None
    version: Optional[str] = None
    scenarios: List["Scenario"] = field(default_factory=list)


@dataclass
class SafetyPreCondition:
    description: str
    expression: Optional[str] = None


@dataclass
class SafetyPostCondition:
    description: str
    expression: Optional[str] = None


# Backwards-compatible aliases used by older modules/tests
SafetyPreConditions = SafetyPreCondition
SafetyPostConditions = SafetyPostCondition


@dataclass
class ResourceBudgets:
    cpu_ms: Optional[int] = None
    memory_mb: Optional[int] = None
    energy_joules: Optional[float] = None
    # compatibility fields
    cpu_cores: Optional[float] = None
    memory_gb: Optional[float] = None
    energy_wh: Optional[float] = None
    time_seconds: Optional[int] = None


@dataclass
class DataSharingPolicy:
    allow_telemetry: bool = True
    allowed_partners: List[str] = field(default_factory=list)
    retention_seconds: Optional[int] = None
    policy_id: Optional[str] = None
    allowed_data: List[str] = field(default_factory=list)


@dataclass
class TelemetrySchema:
    version: str
    signals: Dict[str, str] = field(default_factory=dict)


@dataclass
class Scenario:
    id: Optional[str] = None
    description: Optional[str] = None
    fault_model: Dict[str, Any] = field(default_factory=dict)
    replay_hooks: List[str] = field(default_factory=list)
    scenario_id: Optional[str] = None
    sensors: List[str] = field(default_factory=list)
    environment: Dict[str, Any] = field(default_factory=dict)
    deterministic_replay: bool = False
    seed: Optional[int] = None


@dataclass
class Contract:
    agent: AgentCapabilities
    schema_version: str = "sb-dsl-0.1"
    pre_conditions: List[SafetyPreCondition] = field(default_factory=list)
    post_conditions: List[SafetyPostCondition] = field(default_factory=list)
    resource_budgets: Optional[ResourceBudgets] = None
    data_sharing: Optional[DataSharingPolicy] = None
    telemetry: Optional[TelemetrySchema] = None
    scenario: Optional[Scenario] = None

    def to_json(self) -> str:
        return json.dumps(asdict(self), sort_keys=True, separators=(',', ':'))

    @classmethod
    def from_json(cls, payload: str) -> "Contract":
        d = json.loads(payload)

        def _construct_agent(obj):
            if obj is None:
                return None
            return AgentCapabilities(**obj)

        def _construct_resource(obj):
            return ResourceBudgets(**obj) if obj is not None else None

        def _construct_data_sharing(obj):
            return DataSharingPolicy(**obj) if obj is not None else None

        def _construct_telemetry(obj):
            return TelemetrySchema(**obj) if obj is not None else None

        def _construct_scenario(obj):
            return Scenario(**obj) if obj is not None else None

        pre = [SafetyPreCondition(**pc) for pc in d.get('pre_conditions', [])]
        post = [SafetyPostCondition(**pc) for pc in d.get('post_conditions', [])]

        return cls(
            agent=_construct_agent(d['agent']),
            schema_version=d.get('schema_version', 'sb-dsl-0.1'),
            pre_conditions=pre,
            post_conditions=post,
            resource_budgets=_construct_resource(d.get('resource_budgets')),
            data_sharing=_construct_data_sharing(d.get('data_sharing')),
            telemetry=_construct_telemetry(d.get('telemetry')),
            scenario=_construct_scenario(d.get('scenario')),
        )


class SchemaRegistry:
    def __init__(self):
        self._store: Dict[str, Dict[str, Any]] = {}

    def register(self, name: str, version: str, contract: Contract) -> str:
        key = f"{name}:{version}"
        self._store[key] = {
            "contract": json.loads(contract.to_json()),
            "fingerprint": manifest_fingerprint(contract),
        }
        return key

    def get(self, name: str, version: str) -> Optional[Dict[str, Any]]:
        return self._store.get(f"{name}:{version}")

    def list_entries(self) -> List[str]:
        return list(self._store.keys())


def manifest_fingerprint(contract: Contract) -> str:
    payload = contract.to_json().encode('utf-8')
    return hashlib.sha256(payload).hexdigest()