import json
import logging

from flask import Blueprint, Response, request, stream_with_context

from backends.editor.commands.graph_reset import Params
from backends.editor.commands.graph_reset import reset_graph as reset_graph_cmd
from backends.editor.utils.api import use_db_handler, use_graph_metadata_store
from backends.utils import return_ok
from graph.connections.db_handler import DbHandler
from stores.graph_metadata_store import GraphMetadataStore

graph_instance = Blueprint("graph", __name__, url_prefix="/graph")

logger = logging.getLogger(__name__)


@graph_instance.route("/reset", methods=["POST"])
@use_graph_metadata_store
@use_db_handler(id_getter=lambda: request.get_json().get("graph_id"))
def reset_graph(graph_store: GraphMetadataStore, db_handler: DbHandler) -> Response:
    params = Params(**request.get_json())

    # Check if client accepts SSE
    accept_header = request.headers.get("Accept", "")
    if "text/event-stream" in accept_header:
        # Return SSE stream
        def generate_progress():
            try:
                progress_generator = reset_graph_cmd(graph_store, db_handler, params)

                for progress in progress_generator:
                    # Format as SSE event
                    data = {
                        "stage": progress["stage"],
                        "current_step": progress["current_step"],
                        "total_steps": progress["total_steps"],
                        "current_definition": progress["current_definition"],
                        "percentage": progress["percentage"],
                    }
                    yield f"data: {json.dumps(data)}\n\n"
                    yield ": flush\n\n"

                yield f"data: {json.dumps({'status': 'completed'})}\n\n"

            except Exception as e:
                logger.exception("Error while resetting graph.", exc_info=e)
                # Send error event
                error_data = {"status": "error", "message": str(e)}
                yield f"data: {json.dumps(error_data)}\n\n"

        return Response(
            stream_with_context(generate_progress()),
            mimetype="text/event-stream",
            headers={
                "Cache-Control": "no-cache",
                "Connection": "keep-alive",
                "Access-Control-Allow-Origin": "*",
                "Access-Control-Allow-Headers": "Cache-Control",
                "X-Accel-Buffering": "no",  # for nginx
            },
        )
    else:
        # Traditional non-streaming response - consume the generator
        progress_generator = reset_graph_cmd(graph_store, db_handler, params)

        for _ in progress_generator:
            pass

        return return_ok()
