from __future__ import annotations

from typing import Generator

from pydantic import BaseModel

from graph.connections.db_handler import DbHandler
from graph.ingestion.graph_builder import GraphBuilder, ProgressInfo
from solutions.graph.models import GraphId
from stores.graph_metadata_store import GraphMetadataStore


class Params(BaseModel):
    graph_id: GraphId


def reset_graph(
    store: GraphMetadataStore, db_handler: DbHandler, params: Params
) -> Generator[ProgressInfo, None, None]:
    """
    Raises:
        GraphDoesNotExistError
        GraphMetadataStoreError
        GraphDBWriteInProgressError
        RuntimeError
    """
    graph_id = params.graph_id

    graph_metadata = store.get(graph_id)

    if db_handler.db_exists():
        db_handler.delete_db()
    
    db_handler.create_db()

    with (
        db_handler.get_new_db_instance(readonly=False) as local_db_instance,
        GraphBuilder(local_db_instance) as builder,
    ):
        for progress in builder.insert_all(graph_metadata, track_progress=True):
            yield progress
