from __future__ import annotations

import logging
from typing import Tuple

from pydantic import BaseModel
from typing_extensions import Literal

from backends.editor.connections.db_handler_registry import get_db_handler_registry
from backends.editor.utils.presets import get_webapp_neo4j_preset
from solutions.graph.models import (
    ConnectionConfiguration,
    GraphId,
    ModelValidationError,
    PresetNotSetException,
    Sampling,
    VersionedGraphMetadata,
    validate_sampling,
)
from stores.graph_metadata_store import GraphMetadataStore

logger = logging.getLogger(__name__)

class UpdateConnectionAndSamplingParams(BaseModel):
    graph_id: GraphId
    sampling: Literal["head", "random", "all"] = "head"
    connection_type: Literal["neo4j", "kuzu"]
    max_rows: int = 1000
    ref_version_token: str


def update_connection_and_sampling(store: GraphMetadataStore, params: UpdateConnectionAndSamplingParams) -> Tuple[bool, VersionedGraphMetadata]:
    """
    Updates the connection and sampling configuration for a given graph.

    Args:
        store (GraphMetadataStore): The graph metadata store.
        params (UpdateConnBaseModel): Parameters for updating the connection.

    Returns:
        Tuple[bool, VersionedGraphMetadata]: A tuple containing:
            - bool: True if the connection or the sampling was updated, False otherwise.
            - VersionedGraphMetadata: The updated graph metadata.

    Raises:
    
    Raises:
        ModelValidationError: If the specified preset is not found or is readonly, or the sampling configuration is invalid.
        GraphDoesNotExistError
        ConcurrentUpdateCollisionError
        GraphMetadataStoreError
    """

    graph_id = params.graph_id
    logger.info(f"Starting connection and sample update for graph '{graph_id}'")
    graph_metadata = store.get(graph_id)

    current_connection = graph_metadata["connection"]
    logger.debug(f"Current connection: type='{current_connection['type']}'")

    # Validate connection preset if provided
    if params.connection_type == "neo4j":
        ensure_valid_neo4j_preset(graph_id)

    sampling_updated = (
        graph_metadata["sampling"]["max_rows"] != params.max_rows
        or graph_metadata["sampling"]["sampling"] != params.sampling
    )
    connection_type_updated = graph_metadata["connection"]["type"] != params.connection_type
    
    updated = sampling_updated or connection_type_updated

    if not updated:
        logger.info(f"No connection nor sampling changes detected for graph '{graph_id}' - skipping update")
        return (False, graph_metadata)

    #Validate sampling
    if sampling_updated:
        sampling = Sampling(sampling=params.sampling, max_rows=params.max_rows)
        validate_sampling(sampling)
        graph_metadata["sampling"] = sampling
    
    if connection_type_updated:
        delete_existing_db_and_clear_cache(params, graph_metadata, graph_id)

    # Update metadata in store
    logger.debug(f"Updating graph metadata in store for graph '{graph_id}'")
    
    new_graph_metadata = store.update(graph_metadata, params.ref_version_token)

    logger.info(f"Successful update for graph '{graph_id}'")
    
    return (updated, new_graph_metadata)


def ensure_valid_neo4j_preset(graph_id) -> None:
    preset = get_webapp_neo4j_preset()
    if not preset:
        logger.warning(f"Connection update failed for graph '{graph_id}': Preset not found")
        raise ModelValidationError(f"Preset not found")
    if preset.readonly:
        raise ModelValidationError(f"Preset is readonly")


def delete_existing_db_and_clear_cache(params: UpdateConnectionAndSamplingParams, graph_metadata: VersionedGraphMetadata, graph_id: str) -> None:
    old_connection = graph_metadata["connection"]
    logger.debug(f"Old connection: type='{old_connection['type']}'")

    new_connection = ConnectionConfiguration(type=params.connection_type)
    graph_metadata["connection"] = new_connection
    logger.info(f"Connection changes detected for graph '{graph_id}': type_changed={True}")

    try:
        # Delete existing database if present
        db_handler = get_db_handler_registry().get_handler(graph_id, old_connection["type"])
        if db_handler.db_exists():
            logger.info(f"Deleting existing database for graph '{graph_id}' due to connection change")
            db_handler.delete_db()
        else:
            logger.debug(f"No existing database found for graph '{graph_id}'")
    except PresetNotSetException as e:
        logger.error(f"Failed to get database handler for graph '{graph_id}': {str(e)}", exc_info=True)
    except Exception as e:
        logger.error(f"Failed to cleanup old database for graph '{graph_id}': {str(e)}", exc_info=True)
        raise

    
    # Clear cache to take in account the change of db
    logger.debug(f"Clearing database handler cache for graph '{graph_id}'")
    get_db_handler_registry().clear_cache()