from __future__ import annotations

import logging

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

from editor.backend.utils.api import use_graph_snapshot_store
from editor.backend.utils.webapp_config import webapp_config
from solutions.backend.utils import return_ok
from solutions.graph.models import (
    FeatureUnavailableError,
    ModelValidationError,
    SnapshotDoesNotExistError,
)
from solutions.graph.store.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


@graph_publish.route("", methods=["POST"])
@use_graph_snapshot_store
def publish_snapshot(snapshot_store: GraphMetadataSnapshotStore):
    id = request.args.get("snapshotId", None)
    if not id:
        raise ModelValidationError("Expected a snapshot id to publish.")

    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."
        )

    if not id or not (snapshot := snapshot_store.get_by_id(id)):
        raise SnapshotDoesNotExistError(id)

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

    output_folder_name = f"{snapshot['name']}-{id}"

    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-{id}"
    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": id}
    recipe_settings.save()

    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.")

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

    return return_ok({"recipeName": new_recipe.name, "outputFolderName": output_folder_name, "jobName": job.id})
