"""
SQL generation helpers.

Only shared LLM utilities used by agent tools live here.
"""

import logging

from semantic_models.exceptions import ExecutionError

logger = logging.getLogger(__name__)


# =============================================================================
# LLM Interaction
# =============================================================================


def clean_sql_response(sql: str) -> str:
    """Clean up LLM SQL response (remove markdown code blocks)."""
    if sql.startswith("```"):
        lines = sql.split("\n")
        if lines[0].startswith("```"):
            lines = lines[1:]
        if lines and lines[-1].strip() == "```":
            lines = lines[:-1]
        sql = "\n".join(lines).strip()
    return sql


def call_llm(llm_id: str, prompt: str) -> str:
    text, _ = call_llm_with_trace(llm_id, prompt)
    return text


def call_llm_with_trace(llm_id: str, prompt: str) -> tuple[str, dict | None]:
    try:
        import dataiku

        client = dataiku.api_client()
        project = client.get_default_project()

        llm = project.get_llm(llm_id)
        completion = llm.new_completion()
        completion.with_message(prompt)

        response = completion.execute()
        text = response.text.strip()
        trace = getattr(response, "trace", None)
        return text, trace

    except Exception as e:
        logger.exception("LLM call failed")
        raise ExecutionError(f"LLM call failed: {e}")
