from typing import Dict

from graph.connections.db_connection import DbConnection
from graph.connections.db_instance import DbInstance
from graph.exploration.schema.schema_extractor import EdgeTableSchema, GraphSchema, NodeTableSchema


def get_nodes_table_schema(connection: DbConnection) -> Dict[str, NodeTableSchema]:
    nodes: Dict[str, NodeTableSchema] = {}
    all_groups_info_df = connection.execute("CALL show_tables() RETURN type, name;").get_as_df()  # type: ignore
    for _, group_row in all_groups_info_df.iterrows():
        type = group_row["type"]  # type: ignore
        group_label: str = group_row["name"]

        if type != "NODE":
            continue

        props_info_df = connection.execute(f'CALL TABLE_INFO("{group_label}") RETURN *;').get_as_df()  # type: ignore
        nodes[group_label] = {
            "label": group_label,
            "properties": [
                {
                    "name": prop["name"],  # type: ignore
                    "type": prop["type"],  # type: ignore
                    "is_primary_key": prop["primary key"],  # type: ignore
                }
                for _, prop in props_info_df.iterrows()
            ],
        }

    return nodes


def get_node_group_pk_map(connection: DbConnection) -> Dict[str, str]:
    node_schemas = get_nodes_table_schema(connection)
    result: Dict[str, str] = {}

    for node_group, node_schema in node_schemas.items():
        primary_key_props = [prop for prop in node_schema["properties"] if prop["is_primary_key"]]
        if len(primary_key_props) != 1:
            raise ValueError(f"Expected a single primary key for Kuzu node group '{node_group}', got {primary_key_props}.")
        result[node_group] = primary_key_props[0]["name"]

    return result


def extract_schema(db_instance: DbInstance) -> GraphSchema:
    with db_instance.get_new_conn() as connection:
        nodes = get_nodes_table_schema(connection)
        edges: Dict[str, EdgeTableSchema] = {}

        all_groups_info_df = connection.execute("CALL show_tables() RETURN type, name;").get_as_df()  # type: ignore
        for _, group_row in all_groups_info_df.iterrows():
            type = group_row["type"]  # type: ignore
            group_label: str = group_row["name"]

            if type == "NODE":
                continue
            if type == "REL":
                from_to_info_df = connection.execute(f'CALL SHOW_CONNECTION("{group_label}") RETURN *;').get_as_df()  # type: ignore
                from_group_label: str = ""
                to_group_label: str = ""
                for _, r in from_to_info_df.iterrows():
                    from_group_label = r["source table name"]
                    to_group_label = r["destination table name"]

                props_info_df = connection.execute(f'CALL TABLE_INFO("{group_label}") RETURN *;').get_as_df()  # type: ignore
                edges[group_label] = {
                    "label": group_label,
                    "from_group_label": from_group_label,
                    "to_group_label": to_group_label,
                    "properties": [
                        {
                            "name": prop["name"],  # type: ignore
                            "type": prop["type"],  # type: ignore
                        }
                        for _, prop in props_info_df.iterrows()
                    ],
                }
            else:
                raise Exception(f'Unexpected type encountered "{type}" when extracting Kuzu graph schema.')

    return {"node_groups": nodes, "edge_groups": edges}
