from typing import Dict, List, TypedDict

from graph.connections.db_instance import DbInstance
from graph.connections.kuzu.db_instance import KuzuDbInstance
from graph.connections.neo4j.db_instance import Neo4jDbInstance


class NodePropertySchema(TypedDict):
    name: str
    type: str
    is_primary_key: bool


class EdgePropertySchema(TypedDict):
    name: str
    type: str


class NodeTableSchema(TypedDict):
    label: str
    properties: List[NodePropertySchema]


class EdgeTableSchema(TypedDict):
    label: str
    from_group_label: str
    to_group_label: str
    properties: List[EdgePropertySchema]


class GraphSchema(TypedDict):
    node_groups: Dict[str, NodeTableSchema]
    edge_groups: Dict[str, EdgeTableSchema]


def extract_schema(db_instance: DbInstance) -> GraphSchema:
    from graph.exploration.schema.kuzu.schema_extractor import extract_schema as extract_schema_kuzu
    from graph.exploration.schema.neo4j.schema_extractor import extract_schema as extract_schema_neo4j

    if isinstance(db_instance, KuzuDbInstance):
        return extract_schema_kuzu(db_instance)
    elif isinstance(db_instance, Neo4jDbInstance):
        return extract_schema_neo4j(db_instance)
    else:
        raise ValueError(f"Unsupported database type: {type(db_instance)}")
