"""Project Knowledge Graph relation model."""

from __future__ import annotations

from typing import Any, Optional

from sqlalchemy import Float, ForeignKey, Integer, VARCHAR, UniqueConstraint
from sqlalchemy.orm import Mapped, mapped_column

from .base import Base
from .types import JSONB


class KGRelation(Base):
    __tablename__ = "kg_relations"
    __table_args__ = (
        UniqueConstraint(
            "tenant_id",
            "relation_type",
            "subject_entity_id",
            "object_entity_id",
            name="uq_kg_relations_tenant_type_subject_object",
        ),
    )

    id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
    relation_type: Mapped[str] = mapped_column(VARCHAR(40), nullable=False, index=True)
    subject_entity_id: Mapped[int] = mapped_column(
        Integer,
        ForeignKey("kg_entities.id", ondelete="CASCADE"),
        nullable=False,
        index=True,
    )
    object_entity_id: Mapped[int] = mapped_column(
        Integer,
        ForeignKey("kg_entities.id", ondelete="CASCADE"),
        nullable=False,
        index=True,
    )
    weight: Mapped[float] = mapped_column(Float, nullable=False, default=1.0)
    payload: Mapped[Optional[Any]] = mapped_column(JSONB, nullable=True)
    source_type: Mapped[str] = mapped_column(VARCHAR(32), nullable=False, default="system", index=True)
    source_ref: Mapped[Optional[str]] = mapped_column(VARCHAR(255), nullable=True)