"""IR schema exporter and sample payloads for conformance tests.

Provides small helpers to extract Pydantic JSON schemas for the
EnergiBridge-style IR models and to produce small example payloads.
"""
from __future__ import annotations

from typing import Dict, Any
import json

from . import ir


def get_ir_schemas() -> Dict[str, Any]:
    """Return a mapping of IR model names to their JSON schema (as dicts)."""
    schemas: Dict[str, Any] = {}
    # Intentionally export a small, curated set of models used in the
    # interoperability layer.
    models = [
        ir.LocalMarketContext,
        ir.AggregatedSignalIR,
        ir.PlanDeltaIR,
        ir.DualVariables,
        ir.PrivacyBudgetEntry,
        ir.AuditLogEntry,
        ir.TimeRound,
        ir.GoCRegistryEntry,
    ]
    for m in models:
        # pydantic BaseModel has .schema()
        schemas[m.__name__] = m.schema()
    return schemas


def sample_payloads() -> Dict[str, Any]:
    """Return small example payloads for a few IR types.

    These are useful for documentation and quick conformance testing.
    """
    now = "2020-01-01T00:00:00Z"
    return {
        "LocalMarketContext": {
            "venue_id": "venue-a",
            "symbol": "ABC",
            "timeframe": "1m",
            "objective": "liquidity",
            "timestamp": now,
        },
        "AggregatedSignalIR": {
            "version": "v1",
            "venues": ["venue-a", "venue-b"],
            "feature_vector": {"liquidity_proxy": 0.75},
            "privacy_budget": 0.1,
        },
        "PlanDeltaIR": {
            "contract_id": "c1",
            "author": "agent-1",
            "timestamp": now,
            "actions": {"adjust": "scale", "scale": 0.5},
        },
    }


def schemas_json() -> str:
    """Return the IR schemas as a pretty-printed JSON string."""
    return json.dumps(get_ir_schemas(), indent=2, sort_keys=True)