from uuid import uuid4

from flask import Blueprint, jsonify, request
from pydantic import ValidationError
from werkzeug.exceptions import BadRequest

from backend.config import get_charts_generation_llm
from backend.models.chart import (
    MAX_CHART_SAMPLE_ROWS,
    MAX_CELL_STRING_LENGTH,
    MAX_COLUMN_NAME_LENGTH,
    MAX_COLUMNS,
    MAX_QUERY_LENGTH,
    MAX_REQUEST_ROWS,
    ChartPlanDecisionRequest,
    PlanRequest,
    UpdateChartPlanRequest,
)
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_store

logger = get_logger(__name__)

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


def _find_stored_record(message, req: PlanRequest) -> tuple[list[str], list[list], str | None] | None:
    try:
        artifact_group = message.artifacts_metadata[req.artifacts_id]
        if artifact_group.get("artifacts_id") not in (None, req.artifacts_id):
            return None
        artifact = artifact_group["artifacts"][req.artifact_index]
        records = [part["records"] for part in artifact["parts"] if part.get("type") == "RECORDS"]
        record = records[req.records_index]
        columns, data = record["columns"], record["data"]
    except (AttributeError, IndexError, KeyError, TypeError):
        return None

    if not isinstance(columns, list) or not isinstance(data, list):
        return None
    stored_query = artifact_group.get("query")
    return columns, data, stored_query if isinstance(stored_query, str) else None


def _sample_stored_table(columns: list, data: list) -> list[list] | None:
    if not data or not 1 <= len(columns) <= MAX_COLUMNS:
        return None
    if any(not isinstance(column, str) or not 1 <= len(column) <= MAX_COLUMN_NAME_LENGTH for column in columns):
        return None

    sample = data[:MAX_CHART_SAMPLE_ROWS]
    # TODO: Consider validating cells.
    return sample


@chart_bp.route("/plan", methods=["POST"])
@log_http_request
def chart_plan():
    try:
        req = PlanRequest.model_validate(request.get_json(force=True, silent=False))
    except ValidationError as e:
        details = e.errors(include_input=False, include_context=False)
        logger.warning("Invalid chart request: %s", details)
        return jsonify({"error": "Invalid request", "details": details}), 400
    except BadRequest:
        return jsonify({"error": "Bad JSON"}), 400

    store = get_store()
    message = store.get_message(req.msg_id)
    if not message:
        return jsonify({"error": "Message not found"}), 404

    conversation = store.get_conversation(req.conv_id)
    if not conversation:
        return jsonify({"error": "Conversation not found"}), 404
    actual_msg_index = next(
        (index for index, item in enumerate(conversation.messages) if getattr(item, "id", None) == req.msg_id),
        None,
    )
    if actual_msg_index is None:
        return jsonify({"error": "Message does not belong to the supplied conversation"}), 409

    stored_record = _find_stored_record(message, req)
    if stored_record is None:
        return jsonify({"error": "Artifact record not found"}), 404
    columns, data, stored_query = stored_record

    if req.columns and req.columns != columns:
        return jsonify({"error": "Submitted chart columns do not match the stored artifact record"}), 409
    if req.data and req.data != data[: len(req.data)]:
        logger.warning(
            "Submitted chart data does not match the stored artifact record",
        )
        return jsonify({"error": "Submitted chart data does not match the stored artifact record"}), 409

    sample = _sample_stored_table(columns, data)
    if sample is None:
        logger.warning("Stored artifact record is not valid for chart generation")
        return (
            jsonify(
                {
                    "error": "Stored artifact record is not valid for chart generation",
                    "error_code": "invalid_chart_record",
                }
            ),
            409,
        )

    if stored_query is not None and req.query != stored_query:
        return jsonify({"error": "Query does not match the stored artifact"}), 409

    history = []
    for item in conversation.messages[max(0, actual_msg_index - 20) : actual_msg_index]:
        role = getattr(item, "role", None)
        content = getattr(item, "content", None)
        if role in {"system", "user", "assistant"} and isinstance(content, str):
            history.append({"role": role, "content": content[:MAX_QUERY_LENGTH]})
    if not history:
        return jsonify({"error": "Conversation has no usable chart-generation history"}), 409
    history[-1]["content"] = stored_query or req.query

    try:
        plan = VisualizationService.prepare_chart_plan(
            llm_id=get_charts_generation_llm(),
            messages=history,
            data=sample,
            columns=columns,
            previous_chart_plan=req.previous_chart_plan.model_dump() if req.previous_chart_plan else None,
            user_instructions=req.user_instructions,
            row_count=len(data),
        )
    except Exception:
        logger.exception("Failed to create chart plan")
        plan = None

    generation_failed = not isinstance(plan, dict) or (
        plan.get("chart_type") == "none" and not plan.get("justification")
    )
    if generation_failed:
        logger.error("Failed to create chart plan")
        instructions_hint = (
            " You can try regenerating it with AI and add guidance." if req.user_instructions is None else ""
        )
        plan = {
            "chart_type": "none",
            "justification": f"Chart generation failed for this artifact.{instructions_hint}",
            "chart_data": None,
        }

    plan_entry = {
        "chart_plan": plan,
        "artifacts_id": req.artifacts_id,
        "source_name": req.agent_name,
        "records_index": req.records_index,
        "artifact_index": req.artifact_index,
        "chart_plan_index": len((message.actions or {}).get("chart_plans") or []),
        "generation_failed": generation_failed,
    }
    if req.persist:
        committed_plan = store.append_chart_plan(req.msg_id, plan_entry)
        if committed_plan is None:
            return jsonify({"error": "Message not found"}), 404
        return jsonify(committed_plan), 200

    plan_entry["proposal_id"] = str(uuid4())
    return jsonify(plan_entry), 200


@chart_bp.route("/plan/decision", methods=["POST"])
@log_http_request
def decide_chart_plan():
    """Persist a previously generated AI-edit proposal only after explicit user choice."""
    try:
        req = ChartPlanDecisionRequest.model_validate(request.get_json(force=True, silent=False))
    except ValidationError as e:
        details = e.errors(include_input=False, include_context=False)
        logger.warning("Invalid chart decision request: %s", details)
        return jsonify({"error": "Invalid request", "details": details}), 400
    except BadRequest:
        return jsonify({"error": "Bad JSON"}), 400

    try:
        store = get_store()
        result = store.apply_chart_plan_decision(
            req.message_id,
            chart_plan_index=req.chart_plan_index,
            action=req.action,
            proposal_id=req.proposal_id,
            expected_chart_plan=req.expected_chart_plan.model_dump() if req.expected_chart_plan else None,
            chart_plan=req.chart_plan.model_dump(),
        )
    except Exception:
        logger.exception("Failed to apply chart plan decision")
        return jsonify({"error": "Failed to apply chart plan decision"}), 500

    error = result.get("error")
    if error == "message_not_found":
        return jsonify({"error": "Message not found"}), 404
    if error == "chart_plan_not_found":
        return jsonify({"error": "Chart plan not found"}), 404
    if error in {"chart_plan_changed", "invalid_previous_decision"}:
        return jsonify({"error": "Chart changed before the decision was applied"}), 409
    return jsonify(result), 200


@chart_bp.route("/plan", methods=["PUT"])
@log_http_request
def update_chart_plan():
    try:
        req = UpdateChartPlanRequest.model_validate(request.get_json(force=True, silent=False))
    except ValidationError as e:
        details = e.errors(include_input=False, include_context=False)
        logger.warning("Invalid chart update request: %s", details)
        return jsonify({"error": "Invalid request", "details": details}), 400
    except BadRequest:
        return jsonify({"error": "Bad JSON"}), 400

    try:
        store = get_store()
        result = store.update_chart_plan(
            req.message_id,
            chart_plan_index=req.chart_plan_index,
            chart_data=req.chart_data.model_dump(),
        )
    except Exception:
        logger.exception("Failed to update chart plan")
        return jsonify({"error": "Failed to update chart plan"}), 500

    error = result.get("error")
    if error == "message_not_found":
        return jsonify({"error": "Message not found"}), 404
    if error == "chart_plan_not_found":
        return jsonify({"error": "Chart plan not found"}), 404
    if error == "chart_plan_payload_missing":
        return jsonify({"error": "Chart plan payload missing"}), 400
    return jsonify(result), 200
