from __future__ import annotations

import logging
from typing import Optional

import dataiku
from dataikuapi.dss.project import DSSProject
from dataikuapi.dss.recipe import DSSRecipe, DSSRecipeCreator
from flask import Blueprint, request
from pydantic import BaseModel, Field

from backends.editor.utils.api import use_graph_snapshot_store
from backends.editor.utils.webapp_config import webapp_config
from backends.utils import return_ok
from solutions.graph.models import (
    DbType,
    FeatureUnavailableError,
    SnapshotDoesNotExistError,
    create_neo4j_db_name,
    validate_neo4j_db_name,
)
from stores.graph_metadata_snapshot_store import GraphMetadataSnapshotStore

logger = logging.getLogger(__name__)

graph_publish = Blueprint("graph_publish", __name__, url_prefix="/publish")


class CreateSnapshotParams(BaseModel):
    name: str
    comment: str


class PublishRequestBody(BaseModel):
    snapshotId: str = Field(min_length=1)
    databaseType: DbType = "kuzu"  # 'kuzu' | 'neo4j'
    databaseName: Optional[str] = None  # Optional for Neo4j


@graph_publish.route("", methods=["POST"])
@use_graph_snapshot_store
def publish_snapshot(snapshot_store: GraphMetadataSnapshotStore):
    body = PublishRequestBody(**request.get_json())

    if not snapshot_store:
        raise FeatureUnavailableError(
            "Publish saved configuration feature is not available. You need to configure the Saved configurations dataset."
        )

    output_folder_connection = webapp_config.build_graph_recipe_output_connection
    if not output_folder_connection:
        raise FeatureUnavailableError(
            "Publish saved configuration feature is not available. You need to configure an output folder connection."
        )

    snapshot = snapshot_store.get_by_id(body.snapshotId)
    if not snapshot:
        raise SnapshotDoesNotExistError(body.snapshotId)

    project = DSSProject(dataiku.api_client(), dataiku.default_project_key())

    database_name: str | None = None
    if body.databaseType == "neo4j":
        database_name = body.databaseName.strip() if body.databaseName else None
        if not database_name:
            database_name = create_neo4j_db_name("dku", str(dataiku.default_project_key()), body.snapshotId)
        validate_neo4j_db_name(database_name)

    output_folder_name = f"{snapshot['name']}-{body.snapshotId}"

    folder_type = dataiku.api_client().get_connection(output_folder_connection).get_info().get_type()

    logger.debug(
        f"Creating managed folder '{output_folder_name}' on connection '{output_folder_connection}' of type '{folder_type}'..."
    )
    output_folder = project.create_managed_folder(
        output_folder_name, folder_type=None, connection_name=output_folder_connection
    )
    logger.debug(f"Managed folder '{output_folder_name}' (odb_id:'{output_folder.odb_id}') created.")

    recipe_name = f"build-graph-{body.snapshotId}"
    logger.debug(f"Creating recipe '{recipe_name}'...")

    recipe_creator = DSSRecipeCreator(
        "CustomCode_visual-graph-build-graph",
        name=recipe_name,
        project=project,
    )

    recipe_creator.with_input(webapp_config.snapshots_ds, role="snapshots_ds")  # type: ignore

    # Input datasets should appear only once or Dataiku misbehaves: recipe created but not linked to output and not buildable.
    input_datasets = set()
    for node_groups in snapshot["nodes"].values():
        for node_group_definition in node_groups["definitions"]:
            input_datasets.add(node_group_definition["source_dataset"])

    for edge_groups in snapshot["edges"].values():
        for edge_group_definition in edge_groups["definitions"]:
            input_datasets.add(edge_group_definition["edge_dataset"])

    for ds in input_datasets:
        recipe_creator.with_input(ds, role="nodes_and_edges_datasets_input")  # type: ignore

    recipe_creator.with_output(output_id=output_folder.odb_id, role="main")  # type: ignore

    new_recipe: DSSRecipe = recipe_creator.create()

    logger.debug(f"Recipe '{recipe_name}' created.")

    recipe_settings = new_recipe.get_settings()

    recipe_settings.raw_params["customConfig"] = {
        "snapshot_id": body.snapshotId,
        "database_type": body.databaseType,  # "kuzu" or "neo4j"
        "database_name": database_name  # Always filled for Neo4j, None for Kuzu
    }
    recipe_settings.save()

    output_folder_name = dataiku.Folder(output_folder.id).get_name()

    # For Kuzu (built-in): Create recipe AND auto-start job
    # For Neo4j: Create recipe but DON'T start job (user must configure preset in recipe UI first)
    if body.databaseType == "kuzu":
        job_builder = project.new_job("NON_RECURSIVE_FORCED_BUILD")
        job_builder.with_output(output_folder.odb_id, object_type="MANAGED_FOLDER", object_project_key=project.project_key)
        job = job_builder.start()

        logger.debug(f"Job building '{output_folder_name}' started for Kuzu database.")

        return return_ok({
            "recipeName": new_recipe.name,
            "outputFolderName": output_folder_name,
            "outputFolderId": output_folder.id,
            "projectKey": project.project_key,
            "jobName": job.id
        })
    else:
        # Neo4j: Recipe created but job NOT started
        logger.debug(
            f"Recipe '{recipe_name}' created for Neo4j. "
            f"User must configure Neo4j preset on recipe before running."
        )

        return return_ok({
            "recipeName": new_recipe.name,
            "outputFolderName": output_folder_name,
            "outputFolderId": output_folder.id,
            "projectKey": project.project_key
        })
