from pydantic import BaseModel, Field
from typing import Mapping, Any, List


class Entity(BaseModel):
    """List of entities extracted from the text to form the knowledge graph"""

    name: str = Field(
        description="Name of the entity, it should be a clear and concise term"
    )
    description: str = Field(
        description=(
            "Description of the entity, it should be a complete and comprehensive sentence, not few words. "
            "Sample description of entity 'TiDB in-place upgrade': "
            "'Upgrade TiDB component binary files to achieve upgrade, generally use rolling upgrade method'"
        )
    )
    metadata: Mapping[str, Any] = Field(
        description=(
            "The covariates (which is a comprehensive json TREE, the first field is always: 'topic') to claim the entity. "
        )
    )


class EntityWithID(Entity):
    """Entity extracted from the text to form the knowledge graph with an ID."""

    id: int = Field(description="Unique identifier for the entity.")


class SynopsisInfo(BaseModel):
    """A synopsis corresponds to a group of entities that share the same topic and can contribute to synopsis topic."""

    topic: str = Field(
        description="The shared topic of the synopsis, and each entity in the group can contribute factual data from its own perspective."
    )
    entities: List[int] = Field(
        description="A group of entity(only IDs) that can contribute to the synopsis base on the analysis of entity descriptions and metadata."
    )


class SynopsisEntity(Entity):
    """Unified synopsis entity with comprehensive description and metadata based on the entities group."""

    group_info: SynopsisInfo = Field(
        description="Group of entities to be unified into a single synopsis entity."
    )


class ExistingSynopsisEntity(SynopsisEntity):
    """Unified synopsis entity with comprehensive description and metadata based on the entities group."""

    id: int = Field(description="Unique identifier for the entity.")


class Relationship(BaseModel):
    """List of relationships extracted from the text to form the knowledge graph"""

    source_entity: str = Field(
        description="Source entity name of the relationship, it should an existing entity in the Entity list"
    )
    target_entity: str = Field(
        description="Target entity name of the relationship, it should an existing entity in the Entity list"
    )
    relationship_desc: str = Field(
        description=(
            "Description of the relationship, it should be a complete and comprehensive sentence, not few words. "
            "Sample relationship description: 'TiDB will release a new LTS version every 6 months.'"
        )
    )


class KnowledgeGraph(BaseModel):
    """Graph representation of the knowledge for text."""

    entities: List[Entity] = Field(
        description="List of entities in the knowledge graph"
    )
    relationships: List[Relationship] = Field(
        description="List of relationships in the knowledge graph"
    )


class EntityCovariateInput(BaseModel):
    """List of entities extracted from the text to form the knowledge graph"""

    name: str = Field(description="Name of the entity")
    description: str = Field(description=("Description of the entity"))


class EntityCovariateOutput(BaseModel):
    """List of entities extracted from the text to form the knowledge graph"""

    name: str = Field(description="Name of the entity")
    description: str = Field(description=("Description of the entity"))
    covariates: Mapping[str, Any] = Field(
        description=(
            "The attributes (which is a comprehensive json TREE, the first field is always: 'topic') to claim the entity. "
        )
    )