from __future__ import annotations

import json
from collections.abc import Callable
from dataclasses import dataclass
from typing import Any, TypedDict, cast

from trazaeo._runtime import load_trazaeo_attrs

from trazaeo_workflows.errors import ConversionRuntimeError, DependencyUnavailableError
from trazaeo_workflows.runtime_payloads import (
    load_json_object,
    require_object_field,
    require_optional_str_field,
    require_str_field,
    require_str_list_field,
)
from trazaeo_workflows.workflow_types import VerificationMode

DatasetBundleBuilder = Callable[[str], str]
DeliveryProofPackageBuilder = Callable[[str, str], str]
FullRootProofPackageBuilder = Callable[[str, int, int | None], str]
DeliveryProofReportVerifier = Callable[[str, str | None, str | None], str]


@dataclass(frozen=True)
class DatasetWorkflowBoundary:
    build_dataset_bootstrap_bundle_json: DatasetBundleBuilder | None = None
    build_dataset_incremental_bundle_json: DatasetBundleBuilder | None = None
    build_dataset_delivery_proof_package_json: DeliveryProofPackageBuilder | None = None
    full_root_proof_package_json: FullRootProofPackageBuilder | None = None
    verify_delivery_proof_package_report_json: DeliveryProofReportVerifier | None = None


class SourceFileEntry(TypedDict):
    source_uri: str
    content_hash: str
    byte_length: int
    observed_mtime: str | None


class DatasetPublishedArtifactInput(TypedDict):
    artifact_id: str
    artifact_ref: str
    content_root_hash: str
    media_type: str
    content_descriptor_ref: str | None
    content_descriptor_hash: str | None
    snapshot_id: str | None
    commit_id: str | None


class DatasetWorkflowCommonInput(TypedDict):
    schema_version: str
    issued_at: str
    signed_at: str
    signer_id: str
    signing_key_hex: str
    source_manifest_id: str
    source_dataset_id: str
    source_files: list[SourceFileEntry]
    transform_subject_id: str
    transform_job_id: str
    transform_stage: str
    publish_subject_id: str
    dataset_id: str
    dataset_version: str
    published_artifact: DatasetPublishedArtifactInput
    toolchain: str
    parameters_ref: str
    parameters_hash: str
    determinism_profile: str
    runtime_env_ref: str | None
    runtime_env_hash: str | None
    transform_spec_ref: str
    transform_spec_hash: str
    chunking_profile_ref: str | None
    chunking_profile_hash: str | None
    execution_manifest_ref: str | None
    execution_manifest_hash: str | None
    runtime_manifest_ref: str | None
    runtime_manifest_hash: str | None
    checkpoint_id: str
    checkpoint_time_window: str
    verification_policy_id: str
    stac_refs: list[str]
    reward_context_ref: str | None
    reward_context_hash: str | None
    reward_eligible: bool
    trust_policy: dict[str, Any]
    verification_mode: VerificationMode


class DatasetBootstrapWorkflowInput(TypedDict):
    common: DatasetWorkflowCommonInput
    bootstrap_origin_label: str


class DatasetProvenanceBundle(TypedDict):
    source_manifest: dict[str, Any]
    source_manifest_ref: str
    source_manifest_hash: str
    transform_envelope: dict[str, Any]
    checkpoint_manifest: dict[str, Any]
    checkpoint_manifest_ref: str
    checkpoint_manifest_hash: str
    publish_envelope: dict[str, Any]
    verification_report: dict[str, Any]
    lineage_envelopes: list[str]
    signature_bundle: list[str]
    store_ref: str
    snapshot_id: str | None
    commit_id: str | None


class DatasetIncrementalWorkflowInput(TypedDict):
    common: DatasetWorkflowCommonInput
    prior_bundle: DatasetProvenanceBundle


class DeliveryProofPackage(TypedDict):
    artifact_id: str
    checkpoint_manifest: dict[str, Any]
    lineage_envelopes: list[str]
    signature_bundle: list[str]


class FullRootProofPackage(TypedDict):
    artifact_id: str
    content_descriptor: dict[str, Any]
    content_proof_type: str
    verified_ranges: list[dict[str, Any]]


class DeliveryProofReport(TypedDict):
    package_state: str
    content_state: str
    lineage_state: str
    artifact_id: str
    artifact_path: str | None
    trusted_checkpoint_log_root_hex: str | None


def load_dataset_workflow_boundary(*attr_names: str) -> DatasetWorkflowBoundary:
    runtime = load_trazaeo_attrs(
        *attr_names,
        error_cls=DependencyUnavailableError,
    )
    return DatasetWorkflowBoundary(
        build_dataset_bootstrap_bundle_json=cast(
            DatasetBundleBuilder | None,
            runtime.get("build_dataset_bootstrap_bundle_json"),
        ),
        build_dataset_incremental_bundle_json=cast(
            DatasetBundleBuilder | None,
            runtime.get("build_dataset_incremental_bundle_json"),
        ),
        build_dataset_delivery_proof_package_json=cast(
            DeliveryProofPackageBuilder | None,
            runtime.get("build_dataset_delivery_proof_package_json"),
        ),
        full_root_proof_package_json=cast(
            FullRootProofPackageBuilder | None,
            runtime.get("full_root_proof_package_json"),
        ),
        verify_delivery_proof_package_report_json=cast(
            DeliveryProofReportVerifier | None,
            runtime.get("verify_delivery_proof_package_report_json"),
        ),
    )


def _call_runtime_json(
    operation: str,
    action: Callable[[], str],
) -> dict[str, Any]:
    try:
        payload_json = action()
    except (TypeError, ValueError, KeyError) as exc:
        raise ConversionRuntimeError(f"{operation} failed: {exc}") from exc
    return load_json_object(
        payload_json,
        operation=operation,
        error_cls=ConversionRuntimeError,
    )


def _validate_dataset_provenance_bundle(
    payload: dict[str, Any],
    *,
    operation: str,
) -> DatasetProvenanceBundle:
    return {
        "source_manifest": require_object_field(
            payload,
            key="source_manifest",
            operation=operation,
            error_cls=ConversionRuntimeError,
        ),
        "source_manifest_ref": require_str_field(
            payload,
            key="source_manifest_ref",
            operation=operation,
            error_cls=ConversionRuntimeError,
        ),
        "source_manifest_hash": require_str_field(
            payload,
            key="source_manifest_hash",
            operation=operation,
            error_cls=ConversionRuntimeError,
        ),
        "transform_envelope": require_object_field(
            payload,
            key="transform_envelope",
            operation=operation,
            error_cls=ConversionRuntimeError,
        ),
        "checkpoint_manifest": require_object_field(
            payload,
            key="checkpoint_manifest",
            operation=operation,
            error_cls=ConversionRuntimeError,
        ),
        "checkpoint_manifest_ref": require_str_field(
            payload,
            key="checkpoint_manifest_ref",
            operation=operation,
            error_cls=ConversionRuntimeError,
        ),
        "checkpoint_manifest_hash": require_str_field(
            payload,
            key="checkpoint_manifest_hash",
            operation=operation,
            error_cls=ConversionRuntimeError,
        ),
        "publish_envelope": require_object_field(
            payload,
            key="publish_envelope",
            operation=operation,
            error_cls=ConversionRuntimeError,
        ),
        "verification_report": require_object_field(
            payload,
            key="verification_report",
            operation=operation,
            error_cls=ConversionRuntimeError,
        ),
        "lineage_envelopes": require_str_list_field(
            payload,
            key="lineage_envelopes",
            operation=operation,
            error_cls=ConversionRuntimeError,
        ),
        "signature_bundle": require_str_list_field(
            payload,
            key="signature_bundle",
            operation=operation,
            error_cls=ConversionRuntimeError,
        ),
        "store_ref": require_str_field(
            payload,
            key="store_ref",
            operation=operation,
            error_cls=ConversionRuntimeError,
        ),
        "snapshot_id": require_optional_str_field(
            payload,
            key="snapshot_id",
            operation=operation,
            error_cls=ConversionRuntimeError,
        ),
        "commit_id": require_optional_str_field(
            payload,
            key="commit_id",
            operation=operation,
            error_cls=ConversionRuntimeError,
        ),
    }


def _validate_delivery_proof_package(
    payload: dict[str, Any],
    *,
    operation: str,
) -> DeliveryProofPackage:
    return {
        "artifact_id": require_str_field(
            payload,
            key="artifact_id",
            operation=operation,
            error_cls=ConversionRuntimeError,
        ),
        "checkpoint_manifest": require_object_field(
            payload,
            key="checkpoint_manifest",
            operation=operation,
            error_cls=ConversionRuntimeError,
        ),
        "lineage_envelopes": require_str_list_field(
            payload,
            key="lineage_envelopes",
            operation=operation,
            error_cls=ConversionRuntimeError,
        ),
        "signature_bundle": require_str_list_field(
            payload,
            key="signature_bundle",
            operation=operation,
            error_cls=ConversionRuntimeError,
        ),
    }


def _validate_full_root_proof_package(
    payload: dict[str, Any],
    *,
    operation: str,
) -> FullRootProofPackage:
    verified_ranges = payload.get("verified_ranges")
    if not isinstance(verified_ranges, list) or any(
        not isinstance(entry, dict) for entry in verified_ranges
    ):
        raise ConversionRuntimeError(
            f"{operation} omitted required list[object] field: verified_ranges"
        )
    return {
        "artifact_id": require_str_field(
            payload,
            key="artifact_id",
            operation=operation,
            error_cls=ConversionRuntimeError,
        ),
        "content_descriptor": require_object_field(
            payload,
            key="content_descriptor",
            operation=operation,
            error_cls=ConversionRuntimeError,
        ),
        "content_proof_type": require_str_field(
            payload,
            key="content_proof_type",
            operation=operation,
            error_cls=ConversionRuntimeError,
        ),
        "verified_ranges": cast(list[dict[str, Any]], verified_ranges),
    }


def _validate_delivery_proof_report(
    payload: dict[str, Any],
    *,
    operation: str,
) -> DeliveryProofReport:
    return {
        "package_state": require_str_field(
            payload,
            key="package_state",
            operation=operation,
            error_cls=ConversionRuntimeError,
        ),
        "content_state": require_str_field(
            payload,
            key="content_state",
            operation=operation,
            error_cls=ConversionRuntimeError,
        ),
        "lineage_state": require_str_field(
            payload,
            key="lineage_state",
            operation=operation,
            error_cls=ConversionRuntimeError,
        ),
        "artifact_id": require_str_field(
            payload,
            key="artifact_id",
            operation=operation,
            error_cls=ConversionRuntimeError,
        ),
        "artifact_path": require_optional_str_field(
            payload,
            key="artifact_path",
            operation=operation,
            error_cls=ConversionRuntimeError,
        ),
        "trusted_checkpoint_log_root_hex": require_optional_str_field(
            payload,
            key="trusted_checkpoint_log_root_hex",
            operation=operation,
            error_cls=ConversionRuntimeError,
        ),
    }


def build_dataset_bootstrap_bundle(
    input_payload: DatasetBootstrapWorkflowInput,
) -> DatasetProvenanceBundle:
    boundary = load_dataset_workflow_boundary("build_dataset_bootstrap_bundle_json")
    return _validate_dataset_provenance_bundle(
        _call_runtime_json(
            "build_dataset_bootstrap_bundle_json",
            lambda: cast(
                DatasetBundleBuilder,
                boundary.build_dataset_bootstrap_bundle_json,
            )(json.dumps(input_payload)),
        ),
        operation="build_dataset_bootstrap_bundle_json",
    )


def build_dataset_incremental_bundle(
    input_payload: DatasetIncrementalWorkflowInput,
) -> DatasetProvenanceBundle:
    boundary = load_dataset_workflow_boundary("build_dataset_incremental_bundle_json")
    return _validate_dataset_provenance_bundle(
        _call_runtime_json(
            "build_dataset_incremental_bundle_json",
            lambda: cast(
                DatasetBundleBuilder,
                boundary.build_dataset_incremental_bundle_json,
            )(json.dumps(input_payload)),
        ),
        operation="build_dataset_incremental_bundle_json",
    )


def build_dataset_delivery_proof_package(
    bundle: DatasetProvenanceBundle,
    range_proof_package: dict[str, Any],
) -> DeliveryProofPackage:
    boundary = load_dataset_workflow_boundary("build_dataset_delivery_proof_package_json")
    return _validate_delivery_proof_package(
        _call_runtime_json(
            "build_dataset_delivery_proof_package_json",
            lambda: cast(
                DeliveryProofPackageBuilder,
                boundary.build_dataset_delivery_proof_package_json,
            )(
                json.dumps(bundle),
                json.dumps(range_proof_package),
            ),
        ),
        operation="build_dataset_delivery_proof_package_json",
    )


def build_local_artifact_full_root_proof_package(
    artifact_path: str,
    *,
    chunk_size: int = 1 << 20,
    threads: int | None = None,
) -> FullRootProofPackage:
    boundary = load_dataset_workflow_boundary("full_root_proof_package_json")
    return _validate_full_root_proof_package(
        _call_runtime_json(
            "full_root_proof_package_json",
            lambda: cast(
                FullRootProofPackageBuilder,
                boundary.full_root_proof_package_json,
            )(artifact_path, chunk_size, threads),
        ),
        operation="full_root_proof_package_json",
    )


def verify_dataset_delivery_proof_report(
    delivery_proof_package: DeliveryProofPackage,
    *,
    artifact_path: str | None = None,
    trusted_checkpoint_log_root_hex: str | None = None,
) -> DeliveryProofReport:
    boundary = load_dataset_workflow_boundary("verify_delivery_proof_package_report_json")
    return _validate_delivery_proof_report(
        _call_runtime_json(
            "verify_delivery_proof_package_report_json",
            lambda: cast(
                DeliveryProofReportVerifier,
                boundary.verify_delivery_proof_package_report_json,
            )(
                json.dumps(delivery_proof_package),
                artifact_path,
                trusted_checkpoint_log_root_hex,
            ),
        ),
        operation="verify_delivery_proof_package_report_json",
    )
