from flask import Blueprint, jsonify, request
from pydantic import ValidationError

from backend.models.chart import PlanRequest
from backend.services.conversation_service import ConversationService
from backend.services.visualization_service import VisualizationService
from backend.utils.logger_utils import log_http_request
from backend.utils.logging_utils import get_logger
from backend.utils.utils import get_charts_generation_llm, get_store

logger = get_logger(__name__)

chart_bp = Blueprint("chart", __name__, url_prefix="/chart")


@chart_bp.route("/plan", methods=["POST"])
@log_http_request
def chart_plan():
    try:
        store = get_store()
        payload = request.get_json(force=True, silent=False)
        req = PlanRequest.model_validate(payload)
        conv_id = req.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=[])
        # get messages up to the assistant message that triggered the chart request
        msgs = conv["messages"][: req.msg_index]
        # replace the user message related to the chart request with the one used for artifacts
        msgs[-1]["content"] = req.query
    except ValidationError as e:
        return jsonify({"error": "Invalid request", "details": e.errors()}), 400
    except Exception as exp:
        logger.exception(f"Failed to parse request JSON: {exp}")
        return jsonify(
            {
                "error": f"Bad JSON",
            }
        ), 400

    plan: dict | None = None
    llm_id = get_charts_generation_llm()
    try:
        plan = VisualizationService.prepare_chart_plan(llm_id=llm_id, messages=msgs, data=req.data, columns=req.columns)
    except Exception as e:
        logger.exception("Failed to create chart plan {e}")
    if plan:
        msg = store.get_message(payload["msg_id"])
        if "actions" not in msg or not isinstance(msg["actions"], dict):
            msg["actions"] = {}
        if not msg["actions"].get("chart_plans"):
            msg["actions"]["chart_plans"] = []
        full_plan_info = {
            "chart_plan": plan,
            "artifacts_id": req.artifacts_id,
            "source_name": req.agent_name,
            "records_index": req.records_index,
            "artifact_index": req.artifact_index,
        }
        msg["actions"]["chart_plans"].append(full_plan_info)
        store.update_message(
            payload["msg_id"],
            {
                "actions": msg["actions"],
            },
        )
        return jsonify(full_plan_info), 200
    else:
        return jsonify({"error": "Failed to create Chart plan"}), 500
