import dataiku
from flask import Blueprint, jsonify, request

from backend.services.conversation_service import ConversationService
from backend.services.stories_service import StoriesService
from backend.utils.logger_utils import log_http_request
from backend.utils.logging_utils import get_logger
from backend.utils.utils import get_store, get_user_base_llm

logger = get_logger(__name__)

stories_blueprint = Blueprint("stories", __name__, url_prefix="/stories")


@stories_blueprint.route("", methods=["POST"])
@log_http_request
def create_story():
    client = dataiku.api_client()
    settings = request.get_json()
    workspace = StoriesService.get_agent_stories_workspace(settings["agent_id"])
    try:
        # -------- retrieve datasets ids from their names and project
        project_key = settings["agent_id"].split(":")[0]
        agent_project = client.get_project(project_key=project_key)
        # TODO should we create the workspace if it doesnt exist ?
        if not workspace:
            raise Exception("Workspace not found")
        # Retrieve datasets references in the workspace
        references = []
        for obj in workspace.list_objects():
            ob_info = obj.get_raw()
            for ds_name in settings["references"]:
                ds = agent_project.get_dataset(ds_name)
                obj_ref = ob_info.get("reference") or {}
                if obj_ref.get("projectKey") == project_key and obj_ref.get("id") == ds.id:
                    references.append(ob_info.get("id"))
        if not references:
            raise Exception(
                f"No valid dataset references found in the workspace. Make sure the used datasets {settings['references']} exist in the workspace {workspace.workspace_key}."
            )
        store = get_store()

        # -------- retrieve conversation  --------------------------------
        conv_id = settings["conv_id"]
        conv = store.get_conversation(conv_id) or {}
        if not conv:
            logger.error("Conversation doesn't exist in DB")
            conv = ConversationService._blank_conversation(conv_id=conv_id, agent_ids=[settings["agent_id"]])

        # -------- build stories required params  -------------------------
        messages = conv["messages"]
        llm_id = get_user_base_llm(store)
        prompt = StoriesService.generate_stories_prompt(llm_id=llm_id, messages=messages)

        # TODO should we check if user has access?
        data = {
            "story": {
                "references": references,
                "title": conv["title"] or "Generated title to do",
                "description": f"Stories created from {conv['title']} conversation",
            }
        }
        logger.info(f"Calling /workspaces/{workspace.workspace_key}/stories with params: {data}")
        # -------- create stories -------------------------
        # It will fail if empty refs or if user has no access to the workspace
        rep = client._perform_json("POST", "/workspaces/%s/objects" % workspace.workspace_key, body=data)

        # -------- update message in db -------------------
        msg = store.get_message(settings["msg_id"])
        msg["actions"]["stories"]["stories_id"] = rep["id"]
        rep["prompt"] = prompt
        store.update_message(settings["msg_id"], {"actions": msg["actions"]})
        return rep, 200
    except Exception as e:
        logger.exception(f"Failed to create story {e}")
    return jsonify({"error": "Failed to create Dataiku Story"}), 500
