Layer 1: Model

class DatasetModel(Dataset):
    storage: StoreModel | None = None
    """Set storage for external lakehouse"""
    public_url_prefix: HttpUrlStr | None = None
    """Public url prefix for resources"""
    shards: int = DEFAULT_SHARDS
    """Number of entity-id hash shards for the parquet store. ``0`` (default)
    means a single shard; huge datasets should configure ``8`` or more at
    creation for bounded per-partition working sets (e.g.
    ``ensure_dataset("big_leak", shards=8)``). Immutable after first
    write – changing it requires a full rewrite."""

    def get_public_prefix(self) -> str | None:
        if self.public_url_prefix:
            return self.public_url_prefix
        if settings.public_url_prefix:
            return render(settings.public_url_prefix, {"dataset": self.name})

class CatalogModel(Catalog):
    storage: StoreModel | None = None
    """Lakehouse storage base path"""

class File(Stats):
    """File metadata model. Arbitrary data can be stored in `extra`, including
    ftm properties that should be added to the generated Entity"""

    model_config = ConfigDict(extra="allow")

    dataset: str
    """Dataset name"""
    checksum: str
    """SHA256 checksum (often referred to as `content_hash`)"""
    extra: dict[str, Any] = {}
    """Arbitrary extra data"""
    origin: str | None = None
    """Origin stage of this file"""

    @field_validator("checksum")
    @classmethod
    def check_checksum(cls, v: str) -> str:
        return validate_checksum(v)

    @field_validator("store")
    @classmethod
    def hide_store(cls, v: str) -> str:
        # always don't include original store uri
        return "lakehouse://"

    @model_validator(mode="before")
    @classmethod
    def collect_extra_fields(cls, data: Any) -> Any:
        if not isinstance(data, dict):
            return data
        known_fields = set(cls.model_fields.keys())
        known_fields.update(["id", "path", "parent"])  # computed_field
        extra = data.get("extra", {})
        extra_fields = {k: v for k, v in data.items() if k not in known_fields}
        if extra_fields:
            data = {k: v for k, v in data.items() if k in known_fields}
            data["extra"] = {**extra, **extra_fields}
        return data

    def to_entity(self) -> StatementEntity:
        """Make an entity for this File"""
        schema = mime_to_schema(self.mimetype)
        entity = make_entity(
            {"id": self.id, "schema": schema},
            entity_type=StatementEntity,
            default_dataset=self.dataset,
        )
        entity.add("contentHash", self.checksum)
        entity.add("fileName", self.name)
        entity.add("fileSize", self.size)
        entity.add("mimeType", self.mimetype)
        entity.add("parent", self.parent)
        for prop in schema.properties:
            if prop in self.extra:
                entity.add(prop, self.extra[prop])
        return entity

    def make_parents(self) -> StatementEntities:
        """Make parent `Folder` entities"""
        parent = Path(self.key).parent
        if parent.name:
            yield from make_folders(parent, dataset=self.dataset)

    def make_entities(self) -> StatementEntities:
        yield from self.make_parents()
        yield self.to_entity()

    @computed_field
    @property
    def id(self) -> str:
        """The entity id is generated by a hash of the file path and the
        checksum. Uses just the checksum as id if that's the key"""
        if self.key == self.checksum:
            return self.checksum
        return make_file_id(self.key, self.checksum)

    @computed_field
    @property
    def path(self) -> str:
        # "key" can be misleading in the codebase so this is an alias
        return self.key

    @computed_field
    @property
    def parent(self) -> str | None:
        for parent in reversed(list(self.make_parents())):
            return parent.id

    @property
    def blob_path(self) -> str:
        """Relative path to blob in dataset archive"""
        return path.archive_blob(self.checksum)

    @property
    def meta_path(self) -> str:
        """Relative path for this file's metadata json in dataset archive"""
        return path.archive_meta(self.checksum, self.id)

    @classmethod
    def from_info(cls, info: Stats, checksum: str, **data) -> Self:
        data["dataset"] = data.get("dataset", DEFAULT_DATASET)
        data["checksum"] = checksum
        return cls(**{**info.model_dump(), **data})

    def to_document(self) -> Document:
        return Document(
            id=self.id,
            checksum=self.checksum,
            name=self.name,
            path=self.path,
            size=self.size,
            mimetype=self.mimetype,
            updated_at=self.updated_at,
        )

class DatasetMapping(BaseModel):
    """A complete mapping configuration for a dataset file."""

    model_config = ConfigDict(arbitrary_types_allowed=True)

    dataset: str
    content_hash: str
    queries: list[Mapping]

class JobModel(BaseModel):
    """Status model for a (probably long running) job"""

    run_id: str
    started: datetime | None = None
    stopped: datetime | None = None
    last_updated: datetime | None = None
    pending: int = 0
    done: int = 0
    errors: int = 0
    running: bool = False
    exc: str | None = None
    took: timedelta = timedelta()

    @computed_field
    @property
    def name(self) -> str:
        return self.__class__.__name__

    @field_validator("run_id", mode="before")
    @classmethod
    def ensure_run_id(cls, value: str | None = None) -> str:
        """Give a manual run id or create one"""
        return value or ensure_uuid()

    def touch(self) -> None:
        self.last_updated = datetime.now(timezone.utc)

    def stop(self, exc: Exception | None = None) -> None:
        self.running = False
        self.stopped = datetime.now(timezone.utc)
        self.exc = str(exc)
        if self.started and self.stopped:
            self.took = self.stopped - self.started

    @classmethod
    def make(cls, **kwargs) -> Self:
        kwargs["run_id"] = cls.ensure_run_id(kwargs.get("run_id"))
        return cls(**kwargs)

    @cached_property
    def log(self) -> BoundLogger:
        return get_logger(__name__, run_id=self.run_id)

class DatasetJobModel(JobModel):
    """Status model for a (probably long running) job bound to a dataset"""

    dataset: str

    @cached_property
    def log(self) -> BoundLogger:
        return get_logger(
            f"{self.dataset}.{self.name}",
            run_id=self.run_id,
            dataset=self.dataset,
        )