import logging
import math
from typing import Any

from semantic_models.schemas import PlaygroundFullRequest
from flask import Blueprint, jsonify, request

from semantic_models.agent.semantic_model_query import format_tool_response, run_semantic_model_query
from semantic_models.config import override_app_config
from semantic_models.exceptions import ValidationError
from semantic_models.routes.common import get_service

logger = logging.getLogger(__name__)


def _sanitize_for_json(obj: Any) -> Any:
    """Recursively replace NaN/Infinity with None for JSON serialization."""
    if isinstance(obj, float):
        if math.isnan(obj) or math.isinf(obj):
            return None
        return obj
    if isinstance(obj, dict):
        return {k: _sanitize_for_json(v) for k, v in obj.items()}
    if isinstance(obj, list):
        return [_sanitize_for_json(v) for v in obj]
    return obj

bp = Blueprint("playground", __name__, url_prefix="/api/playground")


# =============================================================================
# Routes
# =============================================================================


@bp.route("/execute", methods=["POST"])
def execute_full():
    data = PlaygroundFullRequest.model_validate(request.json)

    if not data.llm_id:
        raise ValidationError("LLM ID is required")

    logger.info("Executing question: %s", data.question[:100])

    overrides: dict[str, Any] = {}
    fields_set = data.model_fields_set

    if "max_rows_per_query" in fields_set and data.max_rows_per_query is not None:
        overrides["max_rows_per_query"] = data.max_rows_per_query
    if "agent_recursion_limit" in fields_set and data.agent_recursion_limit is not None:
        overrides["agent_recursion_limit"] = data.agent_recursion_limit
    if "sources_records_limit" in fields_set and data.sources_records_limit is not None:
        overrides["sources_records_limit"] = data.sources_records_limit
    if "artifacts_records_limit" in fields_set and data.artifacts_records_limit is not None:
        overrides["artifacts_records_limit"] = data.artifacts_records_limit

    service = get_service()

    if overrides:
        with override_app_config(**overrides):
            result = run_semantic_model_query(
                service=service,
                question=data.question,
                semantic_model_id=data.semantic_model_id,
                version_id=data.version_id,
                llm_id=data.llm_id,
                embedding_llm_id=data.embedding_llm_id,
                initialize_indices=False,
            )
    else:
        result = run_semantic_model_query(
            service=service,
            question=data.question,
            semantic_model_id=data.semantic_model_id,
            version_id=data.version_id,
            llm_id=data.llm_id,
            embedding_llm_id=data.embedding_llm_id,
            initialize_indices=False,
        )

    formatted = format_tool_response(
        question=data.question,
        result=result,
        wrap_sources=False,
    )

    trace = result.get("trace", {})

    response_data = {
        "success": result.get("success", False),
        "error": result.get("error"),
        "response": formatted.get("output"),
        "corrections": result.get("corrections", []),
        "glossary_terms_used": result.get("glossary_terms_used", []),
        "sources": formatted.get("sources"),
        "artifacts": formatted.get("artifacts", []),
        "trace": trace,
        "duration_ms": trace.get("total_duration_ms"),
        "sql": result.get("sql"),  # For golden query saving (may be None)
    }

    # Sanitize NaN/Infinity values for valid JSON
    return jsonify(_sanitize_for_json(response_data))
