"""Utility functions for schema export."""

from typing import Any

from .entities import (
    Account,
    Calculation,
    Document,
    Evidence,
    ExpenseItem,
    FormBox,
    IncomeItem,
    Party,
    Payment,
    PropertyAsset,
    Rule,
    TaxpayerProfile,
)
from .requests import (
    DocumentUploadRequest,
    ExtractionRequest,
    FirmSyncRequest,
    HMRCSubmissionRequest,
    RAGSearchRequest,
    ScheduleComputeRequest,
)
from .responses import (
    DocumentUploadResponse,
    ExtractionResponse,
    FirmSyncResponse,
    HMRCSubmissionResponse,
    RAGSearchResponse,
    ScheduleComputeResponse,
)


def get_entity_schemas() -> dict[str, dict[str, Any]]:
    """Export JSON schemas for all models"""
    schemas = {}

    # Core entities
    schemas["TaxpayerProfile"] = TaxpayerProfile.model_json_schema()
    schemas["Document"] = Document.model_json_schema()
    schemas["Evidence"] = Evidence.model_json_schema()
    schemas["IncomeItem"] = IncomeItem.model_json_schema()
    schemas["ExpenseItem"] = ExpenseItem.model_json_schema()
    schemas["Party"] = Party.model_json_schema()
    schemas["Account"] = Account.model_json_schema()
    schemas["PropertyAsset"] = PropertyAsset.model_json_schema()
    schemas["Payment"] = Payment.model_json_schema()
    schemas["Calculation"] = Calculation.model_json_schema()
    schemas["FormBox"] = FormBox.model_json_schema()
    schemas["Rule"] = Rule.model_json_schema()

    # Request/Response models
    schemas["DocumentUploadRequest"] = DocumentUploadRequest.model_json_schema()
    schemas["DocumentUploadResponse"] = DocumentUploadResponse.model_json_schema()
    schemas["ExtractionRequest"] = ExtractionRequest.model_json_schema()
    schemas["ExtractionResponse"] = ExtractionResponse.model_json_schema()
    schemas["RAGSearchRequest"] = RAGSearchRequest.model_json_schema()
    schemas["RAGSearchResponse"] = RAGSearchResponse.model_json_schema()
    schemas["ScheduleComputeRequest"] = ScheduleComputeRequest.model_json_schema()
    schemas["ScheduleComputeResponse"] = ScheduleComputeResponse.model_json_schema()
    schemas["HMRCSubmissionRequest"] = HMRCSubmissionRequest.model_json_schema()
    schemas["HMRCSubmissionResponse"] = HMRCSubmissionResponse.model_json_schema()
    schemas["FirmSyncRequest"] = FirmSyncRequest.model_json_schema()
    schemas["FirmSyncResponse"] = FirmSyncResponse.model_json_schema()

    return schemas