#
#   Imandra Inc.
#
#   speclogician/presentation/models/container.py
#

from __future__ import annotations

from typing import Literal
from pydantic import BaseModel, Field

from speclogician.presentation.models.trace import TestTracePM, LogTracePM
from speclogician.presentation.models.data_artifact import DocRefPM, SrcCodeRefPM


class ArtifactContainerCountsPM(BaseModel):
    # traces
    num_test_traces_total: int = 0
    num_test_traces_matched: int = 0
    num_test_traces_logic_good: int = 0

    num_log_traces_total: int = 0
    num_log_traces_matched: int = 0
    num_log_traces_logic_good: int = 0

    # data artifacts
    num_src_code_arts_total: int = 0
    num_src_code_arts_matched: int = 0

    num_doc_arts_total: int = 0
    num_doc_arts_matched: int = 0


class ArtifactContainerItemsPM(BaseModel):
    test_traces: list[TestTracePM] = Field(default_factory=list)
    log_traces: list[LogTracePM] = Field(default_factory=list)
    doc_refs: list[DocRefPM] = Field(default_factory=list)
    src_code_refs: list[SrcCodeRefPM] = Field(default_factory=list)


class ArtifactContainerPM(BaseModel):
    kind: Literal["artifact_container"] = "artifact_container"
    counts: ArtifactContainerCountsPM
    items: ArtifactContainerItemsPM