try:
    from ProvenaInterfaces.SharedTypes import StatusResponse
    from ProvenaInterfaces.ProvenanceModels import *
    from ProvenaInterfaces.RegistryAPI import ItemSubType
except:
    from .SharedTypes import StatusResponse
    from .ProvenanceModels import *

from typing import Optional, Any, Dict, Type
from pydantic import BaseModel, Field


class ProvenanceRecordInfo(BaseModel):
    # Handle ID
    id: str
    # Prov serialised document
    prov_json: str
    # Identified version of input format
    record: ModelRunRecord


class RegisterModelRunResponse(StatusResponse):
    session_id: Optional[str] = None


class SyncRegisterModelRunResponse(StatusResponse):
    record_info: ProvenanceRecordInfo


class RegisterBatchModelRunRequest(BaseModel):
    records: List[ModelRunRecord]


class RegisterBatchModelRunResponse(StatusResponse):
    session_id: Optional[str] = None


class LineageResponse(StatusResponse):
    record_count: Optional[int] = None
    graph: Optional[Dict[str, Any]] = None


class ConvertModelRunsResponse(StatusResponse):
    # list of model run records
    new_records: Optional[List[ModelRunRecord]] = None

    # list of job ids of validated existing jobs which don't need to be lodged
    existing_records: Optional[List[str]] = None

    warnings: Optional[List[str]] = None


class AddStudyLinkQueryParameters(BaseModel):
    # the model run that is being linked to the study
    model_run_id: str
    # the study that is being linked to the model run
    study_id: str


class AddStudyLinkResponse(StatusResponse):
    # the model run that was linked to the study
    model_run_id: str
    # the study that was linked to the model run
    study_id: str
    # the session id for the provenance graph update job
    session_id: str


class PostUpdateModelRunInput(BaseModel):
    model_run_id: str
    reason: str
    record: ModelRunRecord


class PostUpdateModelRunResponse(BaseModel):
    session_id: str

class GenerateReportRequest(BaseModel):
    id: str
    item_subtype: ItemSubType
    # Greater than or equal to 1 and less than equal to 3
    depth: int = Field(ge=1, le=3)

class PostDeleteGraphRequest(BaseModel):
    record_id: str
    trial_mode: bool = False

class PostDeleteGraphResponse(BaseModel):
    # Diff list - each entry is a diff action with possible other metadata
    diff: List[Dict[str, Any]]
