import copy
import logging
from typing import Any, Optional

from semantic_models.agent import run_agent
from semantic_models.utils.spec_utils import build_spec_from_version
from semantic_models.schemas import SemanticModelSpec
from semantic_models.services.semantic_models_service import SemanticModelsService, Version
from semantic_models.services.sql_execution import get_dialect_from_spec

logger = logging.getLogger(__name__)

def get_version_or_raise(
    service: SemanticModelsService, semantic_model_id: str, version_id: Optional[str]
) -> Version:
    if version_id is not None:
        version = service.get_version(semantic_model_id, version_id)
    else:
        version = service.get_active_version(semantic_model_id)
    if not version:
        raise ValueError(f"No versions found for model {semantic_model_id}")
    return version


def get_dialect_from_spec_safe(spec: SemanticModelSpec) -> str:
    try:
        return get_dialect_from_spec(spec)
    except Exception as e:
        logger.warning("Could not detect dialect from spec: %s", e)
        return "postgres"


def _normalize_embedding_llm_id(value: Any) -> Optional[str]:
    if not isinstance(value, str):
        return None
    trimmed = value.strip()
    return trimmed or None


def initialize_search_indices_from_spec(
    semantic_model_id: str,
    version_id: str,
    spec_dict: dict[str, Any],
    embedding_llm_id: Optional[str],
) -> None:
    from semantic_models.resolution.index import collect_sources_from_spec, initialize_value_indices_from_sources

    sources = collect_sources_from_spec(semantic_model_id, version_id, spec_dict)
    results = initialize_value_indices_from_sources(
        sources=sources,
        embedding_llm_id=embedding_llm_id,
        spec_dict=spec_dict,
    )
    logger.info("Search index init: %s", results)


def run_semantic_model_query(
    service: SemanticModelsService,
    question: str,
    semantic_model_id: str,
    llm_id: str,
    version_id: Optional[str] = None,
    *,
    embedding_llm_id: Optional[str] = None,
    initialize_indices: bool = False,
    trace: Any | None = None,
) -> dict[str, Any]:
    version = get_version_or_raise(service, semantic_model_id, version_id)

    spec = build_spec_from_version(version)
    spec_dict = spec.model_dump()

    if not spec_dict.get("entities"):
        raise ValueError("Semantic model has no entities defined")

    dialect = get_dialect_from_spec_safe(spec)
    logger.info("Using dialect: %s", dialect)

    resolved_embedding_llm_id = _normalize_embedding_llm_id(embedding_llm_id)

    if initialize_indices:
        initialize_search_indices_from_spec(
            semantic_model_id,
            version.id,
            spec_dict,
            resolved_embedding_llm_id,
        )

    return run_agent(
        question=question,
        semantic_model_id=semantic_model_id,
        spec=spec_dict,
        dialect=dialect,
        llm_id=llm_id,
        embedding_llm_id=resolved_embedding_llm_id,
        version_id=version.id,
        trace=trace,
    )


def _default_sources(question: str) -> dict[str, Any]:
    return {
        "toolCallDescription": f"semantic-model-query for: {question[:100]}",
        "items": [],
    }


def format_tool_response(
    question: str,
    result: dict[str, Any],
    *,
    wrap_sources: bool = True,
) -> dict[str, Any]:
    if result.get("success"):
        response = _build_success_response(question, result)
    elif result.get("response") and not result.get("error"):
        response = _build_info_response(question, result)
    else:
        response = _build_error_response(question, result)

    if wrap_sources and "sources" in response and isinstance(response["sources"], dict):
        response["sources"] = [response["sources"]]

    if "trace" in result:
        response["trace"] = result.get("trace")

    return response


def _build_success_response(
    question: str,
    result: dict[str, Any],
) -> dict[str, Any]:
    sources = copy.deepcopy(result.get("sources") or _default_sources(question))
    artifacts = result.get("artifacts", [])
    agent_response = result.get("response", "")
    row_count = result.get("row_count", 0)

    output = agent_response or f"Query returned {row_count} rows."
    response: dict[str, Any] = {"output": output, "sources": sources}
    if artifacts:
        response["artifacts"] = artifacts
    return response


def _build_info_response(question: str, result: dict[str, Any]) -> dict[str, Any]:
    sources = copy.deepcopy(result.get("sources") or _default_sources(question))
    agent_response = result.get("response", "")
    return {"output": agent_response, "sources": sources}


def _build_error_response(question: str, result: dict[str, Any]) -> dict[str, Any]:
    error_msg = result.get("error", "Query failed")
    sources = copy.deepcopy(result.get("sources") or _default_sources(question))

    # Include partial SQL if available
    if result.get("sql"):
        sources["items"].append(
            {
                "type": "GENERATED_SQL_QUERY",
                "performedQuery": result["sql"],
            }
        )

    if not any(item.get("type") == "ERROR" for item in sources.get("items", [])):
        sources["items"].append(
            {
                "type": "ERROR",
                "textSnippet": error_msg,
            }
        )

    return {
        "output": f"Could not answer the question: {error_msg}",
        "sources": sources,
    }


def build_client_error_response(question: str, error: str, *, wrap_sources: bool = True) -> dict[str, Any]:
    sources = {
        "toolCallDescription": f"semantic-model-query for: {question[:100]}",
        "items": [
            {
                "type": "ERROR",
                "textSnippet": error,
            }
        ],
    }

    response = {
        "output": f"Error querying semantic model: {error}",
        "sources": sources,
    }

    if wrap_sources:
        response["sources"] = [sources]
    return response
