import copy
import json
import logging
import time
from dataclasses import dataclass, field
from typing import Annotated, Any, Dict, List, Literal, Optional, TypedDict

from dataiku.llm.tracing import new_trace
from langchain_core.callbacks import Callbacks
from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, SystemMessage, ToolMessage
from langchain_core.runnables import RunnableConfig
from langchain_core.tools import tool
from langgraph.graph import END, StateGraph
from langgraph.graph.message import add_messages
from langgraph.prebuilt import ToolNode

from semantic_models.config import get_app_config
from semantic_models.schemas import SemanticModelSpec

logger = logging.getLogger(__name__)


@dataclass
class RuntimeContext:
    question: str
    spec: dict
    dialect: str
    llm_id: str
    embedding_llm_id: str | None
    semantic_model_id: str
    version_id: str | None

    # Tracking
    trace: Any = None
    trace_root: Any = None
    trace_span: Any = None
    tool_calls_span: Any = None
    sources: Dict[str, Any] = field(default_factory=dict)
    entity_cache: dict[tuple[str, bool], dict[str, Any]] = field(default_factory=dict)
    corrections_cache: list[dict] | None = None
    glossary_cache: list[dict] | None = None
    vocabulary_cache: list[dict] | None = None
    attribute_values_cache: dict[tuple[str, str], dict[str, Any]] = field(default_factory=dict)
    spec_model: SemanticModelSpec = field(init=False)

    # Resolved from AppConfig (internal use only, no per-request overrides)
    _resolved_max_rows: int = field(init=False)
    _resolved_recursion_limit: int = field(init=False)
    _resolved_sources_limit: int = field(init=False)
    _resolved_artifacts_limit: int = field(init=False)
    _resolved_glossary_autoload: bool = field(init=False)
    _resolved_vocabulary_limit: int = field(init=False)

    def __post_init__(self):
        """Resolve values from AppConfig."""
        app = get_app_config()
        self.spec_model = SemanticModelSpec.model_validate(self.spec)
        self._resolved_max_rows = app.max_rows_per_query
        self._resolved_recursion_limit = app.agent_recursion_limit
        self._resolved_sources_limit = app.sources_records_limit
        self._resolved_artifacts_limit = app.artifacts_records_limit
        self._resolved_glossary_autoload = app.glossary_autoload
        self._resolved_vocabulary_limit = app.vocabulary_terms_in_prompt
        if self.embedding_llm_id == "":
            self.embedding_llm_id = None

        # Initialize sources structure
        if not self.sources:
            self.sources = {
                "toolCallDescription": f"semantic-model-query for: {self.question[:100]}",
                "items": [],
            }

    def add_info(self, message: str) -> None:
        """Add an INFO source item."""
        self.sources["items"].append(
            {
                "type": "INFO",
                "textSnippet": message,
            }
        )
        logger.debug("Source added: INFO - %s", message[:100])

    def add_sql_query(self, sql: str) -> None:
        """Add a GENERATED_SQL_QUERY source item."""
        self.sources["items"].append(
            {
                "type": "GENERATED_SQL_QUERY",
                "performedQuery": sql,
            }
        )
        logger.debug("Source added: GENERATED_SQL_QUERY (%d chars)", len(sql))

    def add_records(
        self,
        columns: List[str],
        data: List[List[Any]],
        max_rows: Optional[int] = None,
    ) -> Dict[str, Any]:
        """Add a RECORDS source item with limited rows for LLM context."""
        if max_rows is None:
            max_rows = self._resolved_sources_limit

        truncated_data = data[:max_rows] if max_rows > 0 else data
        truncated = len(data) > len(truncated_data)

        records_item = {
            "type": "RECORDS",
            "records": {
                "columns": columns,
                "data": truncated_data,
            },
        }
        if truncated:
            records_item["truncated"] = True
            records_item["total_rows"] = len(data)

        self.sources["items"].append(records_item)
        logger.debug(
            "Source added: RECORDS (%d cols, %d/%d rows%s)",
            len(columns),
            len(truncated_data),
            len(data),
            " [truncated]" if truncated else "",
        )
        return records_item

    def add_error(self, error: str) -> None:
        """Add an ERROR source item."""
        self.sources["items"].append(
            {
                "type": "ERROR",
                "textSnippet": error,
            }
        )
        logger.warning("Source added: ERROR - %s", error[:200])

    def build_records_artifact(
        self,
        name: str,
        description: str,
        columns: List[str],
        rows: List[List[Any]],
        max_records: Optional[int] = None,
    ) -> Dict[str, Any] | None:
        """Build a RECORDS artifact for agent tool response."""
        if max_records is None:
            max_records = self._resolved_artifacts_limit
        if max_records <= 0:
            return None

        truncated_rows = rows[:max_records] if max_records > 0 else rows
        truncated = len(rows) > len(truncated_rows)

        artifact = {
            "name": name,
            "description": description,
            "parts": [
                {
                    "type": "RECORDS",
                    "records": {
                        "columns": columns,
                        "data": truncated_rows,
                    },
                }
            ],
        }

        if truncated:
            artifact["truncated"] = True
            artifact["total_rows"] = len(rows)

        return artifact


class AgentState(TypedDict):
    question: str
    semantic_model_id: str
    version_id: Optional[str]
    spec: dict
    dialect: str
    llm_id: str

    # Runtime context reference (for accessing resolved config)
    ctx: RuntimeContext

    # LangGraph message history (for ReAct loop)
    messages: Annotated[list[BaseMessage], add_messages]

    # Tracking
    sql_attempts: list[dict]
    sql_execution_count: int
    start_time: float

    # Final output
    final_result: Optional[dict]


def create_tools(ctx: RuntimeContext):
    """Create tools with access to runtime context."""

    spec_raw = ctx.spec
    semantic_model_id = ctx.semantic_model_id
    dialect = ctx.dialect
    llm_id = ctx.llm_id
    spec_model = ctx.spec_model
    resolved_version_id = ctx.version_id
    private_editor_data = None
    golden_queries = []
    service = None
    try:
        from semantic_models.services.factory import get_semantic_models_service

        service = get_semantic_models_service()
    except Exception as e:
        logger.warning("Failed to initialize semantic models service: %s", e)

    if service is not None:
        if resolved_version_id is None:
            try:
                active_version = service.get_active_version(semantic_model_id)
                if active_version:
                    resolved_version_id = active_version.id
            except Exception as e:
                logger.warning("Failed to resolve active version for tools: %s", e)

        if resolved_version_id is not None:
            try:
                version = service.get_version(semantic_model_id, resolved_version_id)
                if version:
                    private_editor_data = version.privateEditorData
                    golden_queries = version.goldenQueries or []
            except Exception as e:
                logger.warning("Failed to load privateEditorData for tools: %s", e)

    def _compact_entity(entity: dict[str, Any]) -> dict[str, Any]:
        compact = dict(entity)
        compact["attributes"] = [
            {
                "name": attr.get("name"),
                "type": attr.get("type"),
                "dssType": attr.get("dssType"),
            }
            for attr in entity.get("attributes", [])
        ]
        compact["metrics"] = [
            {
                "name": metric.get("name"),
                "pseudoSQLExpression": metric.get("pseudoSQLExpression"),
            }
            for metric in entity.get("metrics", [])
        ]
        compact["filters"] = [
            {
                "name": filt.get("name"),
                "pseudoSQLExpression": filt.get("pseudoSQLExpression"),
            }
            for filt in entity.get("filters", [])
        ]
        return compact

    def _get_entity_cached(entity_id: str, compact: bool) -> dict[str, Any] | None:
        cache_key = (entity_id, compact)
        cached = ctx.entity_cache.get(cache_key)
        if cached is not None:
            return cached
        if compact:
            full_key = (entity_id, False)
            full_cached = ctx.entity_cache.get(full_key)
            if full_cached is not None:
                compact_entity = _compact_entity(full_cached)
                ctx.entity_cache[cache_key] = compact_entity
                return compact_entity

        from semantic_models.agent.tools import get_entity_attributes as _get_entity_attributes

        entity_detail = _get_entity_attributes(
            spec_raw,
            entity_id,
            compact=compact,
            semantic_model_id=semantic_model_id,
            version_id=resolved_version_id,
            private_editor_data=private_editor_data,
        )
        if entity_detail is None:
            return None
        ctx.entity_cache[cache_key] = entity_detail
        return entity_detail

    def _get_tool_run_span(callbacks: Callbacks | None) -> Any | None:
        if callbacks is None:
            return None
        try:
            from dataiku.langchain import dku_span_builder_for_callbacks

            span = dku_span_builder_for_callbacks(callbacks, ignore_missing=True)
            if getattr(span, "span", {}).get("name") == "ignored":
                return None
            return span
        except Exception:
            return None

    class _ToolTrace:
        def __init__(self, span: Any, close_span: bool):
            self.span = span
            self.close_span = close_span

        def __getattr__(self, name: str):
            return getattr(self.span, name)

    def _start_tool_trace(
        tool_id: str,
        tool_input: dict[str, Any] | None,
        callbacks: Callbacks | None,
    ) -> Any | None:
        span_name = f"DKU_MANAGED_TOOL_CALL:{tool_id}"
        tool_run_span = _get_tool_run_span(callbacks)
        if tool_run_span is not None:
            try:
                tool_run_span.span["name"] = span_name
            except Exception:
                pass
            tool_run_span.attributes["toolId"] = tool_id
            tool_run_span.attributes["toolName"] = tool_id
            tool_run_span.attributes["toolType"] = f"Custom_agent_tool_semantic_model_query_{tool_id}"
            if tool_input is not None:
                tool_run_span.inputs["input"] = tool_input
            return _ToolTrace(tool_run_span, close_span=False)

        if ctx.tool_calls_span is None:
            return None

        tool_parent = ctx.tool_calls_span.subspan(span_name)
        tool_parent.__enter__()
        tool_parent.attributes["toolId"] = tool_id
        tool_parent.attributes["toolName"] = tool_id
        tool_parent.attributes["toolType"] = f"Custom_agent_tool_semantic_model_query_{tool_id}"
        if tool_input is not None:
            tool_parent.inputs["input"] = tool_input
        return _ToolTrace(tool_parent, close_span=True)

    def _finish_tool_trace(span: Any | None, output: Any = None, error: str | None = None) -> None:
        if span is None:
            return
        tool_span = span
        close_span = False
        if isinstance(span, _ToolTrace):
            tool_span = span.span
            close_span = span.close_span
        if error:
            tool_span.attributes["status"] = "error"
            tool_span.outputs["error"] = error
        if output is not None:
            tool_span.outputs["output"] = output
        if close_span:
            tool_span.__exit__(None, None, None)

    @tool
    def list_entities(callbacks: Callbacks = None) -> str:
        """List all entities in the semantic model with summaries.

        Returns entity names, descriptions, and counts of attributes/metrics/filters.
        Call this first to understand what data is available.
        """
        from semantic_models.agent.tools import list_entities as _list_entities

        span = _start_tool_trace("list_entities", {}, callbacks)
        try:
            result = _list_entities(spec_raw)
            _finish_tool_trace(span, output=result)
            return json.dumps(result, indent=2)
        except Exception as e:
            _finish_tool_trace(span, error=str(e))
            raise

    @tool
    def get_entity_attributes(entity_id: str, compact: bool = False, callbacks: Callbacks = None) -> str:
        """Get detailed attributes for a specific entity.

        Args:
            entity_id: ID of the entity to get attributes for
            compact: If True, return only name and type (for large entities with 100+ columns)

        Returns all columns, metrics, and filters for the entity with their descriptions,
        types, expressions, and sample values.
        """
        span = _start_tool_trace(
            "get_entity_attributes",
            {"entity_id": entity_id, "compact": compact},
            callbacks,
        )
        try:
            result = _get_entity_cached(entity_id, compact)
            if result is None:
                _finish_tool_trace(span, error=f"Entity '{entity_id}' not found")
                return f"Entity '{entity_id}' not found"
            _finish_tool_trace(span, output=result)
            return json.dumps(result, indent=2)
        except Exception as e:
            _finish_tool_trace(span, error=str(e))
            raise

    @tool
    def get_relationships(entity_ids: list[str], callbacks: Callbacks = None) -> str:
        """Get join relationships between entities.

        Args:
            entity_ids: List of entity IDs to get relationships for

        Returns join hints for connecting the specified entities.
        """
        from semantic_models.agent.tools import get_relationships as _get_relationships

        span = _start_tool_trace("get_relationships", {"entity_ids": entity_ids}, callbacks)
        try:
            result = _get_relationships(spec_raw, entity_ids)
            _finish_tool_trace(span, output=result)
            return json.dumps(result, indent=2)
        except Exception as e:
            _finish_tool_trace(span, error=str(e))
            raise

    @tool
    def get_attribute_values(
        entity_id: str,
        attribute: str,
        limit: int = 20,
        callbacks: Callbacks = None,
    ) -> str:
        """Fetch distinct values for one attribute when a filter literal is unclear.

        Args:
            entity_id: Entity ID containing the attribute
            attribute: Attribute/column name
            limit: Maximum values to return (default 20)

        Uses configured distinct-values settings (MANUAL/AUTO_INDEX/NONE).
        """
        from semantic_models.agent.tools import get_attribute_values as _get_attribute_values

        span = _start_tool_trace(
            "get_attribute_values",
            {"entity_id": entity_id, "attribute": attribute, "limit": limit},
            callbacks,
        )
        try:
            result = _get_attribute_values(
                spec_raw,
                entity_id,
                attribute,
                limit=limit,
                semantic_model_id=semantic_model_id,
                version_id=resolved_version_id,
            )

            values = result.get("values") or []
            if values:
                source = result["source"]
                ctx.attribute_values_cache[(entity_id, attribute)] = {
                    "values": values,
                    "source": source,
                }
                logger.info(
                    "Cached targeted values for %s.%s (%d values, source=%s)",
                    entity_id,
                    attribute,
                    len(values),
                    source,
                )

            _finish_tool_trace(span, output=result)
            return json.dumps(result, indent=2)
        except Exception as e:
            _finish_tool_trace(span, error=str(e))
            raise

    @tool
    def resolve_values(question: str, callbacks: Callbacks = None) -> str:
        """Resolve potential data values in the question using fuzzy matching.

        Args:
            question: The user's natural language question

        Extracts terms that might be data values (names, categories, etc.) and
        searches for matches in the indexed column values. Returns corrections
        with the exact values to use in SQL.
        """
        from semantic_models.agent.tools import resolve_values as _resolve_values

        span = _start_tool_trace("resolve_values", {"question": question}, callbacks)
        try:
            if resolved_version_id is None:
                error_msg = "No semantic model version available for resolution"
                payload = {
                    "original_question": question,
                    "resolved_question": question,
                    "corrections": [],
                    "error": error_msg,
                }
                _finish_tool_trace(span, error=error_msg, output=payload)
                return json.dumps(payload, indent=2)

            result = _resolve_values(
                question=question,
                spec=spec_model,
                semantic_model_id=semantic_model_id,
                version_id=resolved_version_id,
                llm_id=llm_id,
                embedding_llm_id=ctx.embedding_llm_id,
                private_editor_data=private_editor_data,
                trace_span=span,
            )
            ctx.corrections_cache = result.get("corrections", [])
            _finish_tool_trace(span, output=result)
            return json.dumps(result, indent=2)
        except Exception as e:
            _finish_tool_trace(span, error=str(e))
            raise

    @tool
    def get_glossary_terms(callbacks: Callbacks = None) -> str:
        """Get business vocabulary terms from the glossary relevant to the question.

        Returns approved glossary terms that map business language to schema elements.
        Only returns terms that appear in the user's question to save context.
        """
        from semantic_models.agent.tools import get_glossary_terms as _get_glossary_terms

        span = _start_tool_trace("get_glossary_terms", {"question": ctx.question, "limit": 50}, callbacks)
        try:
            result = _get_glossary_terms(
                semantic_model_id=semantic_model_id,
                version_id=resolved_version_id,
                spec=spec_model,
                question=ctx.question,
                limit=50,
            )
            ctx.glossary_cache = result
            _finish_tool_trace(span, output=result)
            return json.dumps(result, indent=2)
        except Exception as e:
            _finish_tool_trace(span, error=str(e))
            raise

    @tool
    def generate_sql(
        question: str,
        entity_ids: list[str],
        corrections: Optional[list[dict]] = None,
        glossary_terms: Optional[list[dict]] = None,
        callbacks: Callbacks = None,
    ) -> str:
        """Generate SQL from the question using gathered schema context.

        Args:
            question: The user's natural language question
            entity_ids: List of entity IDs to include in the query context
            corrections: Optional value corrections from resolve_values
            glossary_terms: Optional glossary terms from get_glossary_terms

        Call this after gathering relevant entities, relationships, and value context.
        Returns the generated SQL query.
        """
        from semantic_models.agent.tools import (
            generate_sql as _generate_sql,
        )
        from semantic_models.agent.tools import (
            get_glossary_terms as _get_glossary_terms,
        )
        from semantic_models.agent.tools import (
            get_relationships as _get_relationships,
        )

        span = _start_tool_trace(
            "generate_sql",
            {
                "question": question,
                "entity_ids": entity_ids,
                "corrections": corrections,
                "glossary_terms": glossary_terms,
            },
            callbacks,
        )
        if corrections is not None:
            ctx.corrections_cache = corrections
        if glossary_terms is not None:
            ctx.glossary_cache = glossary_terms

        effective_corrections = corrections
        if effective_corrections is None and ctx.corrections_cache is not None:
            effective_corrections = ctx.corrections_cache

        effective_glossary = glossary_terms
        if effective_glossary is None and ctx.glossary_cache is not None:
            effective_glossary = ctx.glossary_cache

        # Gather entity details
        entities = []
        for eid in entity_ids:
            entity_detail = _get_entity_cached(eid, compact=False)
            if entity_detail:
                entities.append(copy.deepcopy(entity_detail))

        # Inject any targeted values fetched via get_attribute_values
        targeted_values_applied = 0
        for entity in entities:
            entity_name = entity.get("entity_id") or entity.get("name")
            if not entity_name:
                continue
            for attr in entity.get("attributes", []):
                attr_name = attr.get("name")
                if not attr_name:
                    continue
                cached = ctx.attribute_values_cache.get((entity_name, attr_name))
                if not cached:
                    continue
                attr["manualValues"] = list(cached.get("values") or [])
                attr["sampleValuesSource"] = cached["source"]
                targeted_values_applied += 1
        if targeted_values_applied:
            logger.info("Applied targeted sample values to %d attributes for SQL generation", targeted_values_applied)

        # Get relationships between selected entities
        relationships = _get_relationships(spec_raw, entity_ids)

        if effective_glossary is None and ctx._resolved_glossary_autoload:
            version_id_val = resolved_version_id
            if version_id_val is not None:
                logger.info("Auto-loading glossary terms (LLM did not call get_glossary_terms)")
                effective_glossary = _get_glossary_terms(
                    semantic_model_id=semantic_model_id,
                    version_id=version_id_val,
                    spec=spec_model,
                    question=ctx.question,
                    limit=50,
                )
                logger.info("Auto-loaded %d glossary terms", len(effective_glossary) if effective_glossary else 0)
                ctx.glossary_cache = effective_glossary

        # Load vocabulary terms (general business terms not bound to schema)
        effective_vocabulary = ctx.vocabulary_cache
        if effective_vocabulary is None and ctx._resolved_vocabulary_limit > 0:
            version_id_val = resolved_version_id
            if version_id_val is not None:
                from semantic_models.services.glossary_utils import load_vocabulary_terms_for_prompt

                effective_vocabulary = load_vocabulary_terms_for_prompt(
                    semantic_model_id=semantic_model_id,
                    version_id=version_id_val,
                    limit=ctx._resolved_vocabulary_limit,
                )
                logger.info("Loaded %d vocabulary terms", len(effective_vocabulary) if effective_vocabulary else 0)
                ctx.vocabulary_cache = effective_vocabulary

        # Convert golden queries to dicts for prompt
        golden_queries_for_prompt = [
            {"question": gq.question, "sql": gq.generatedSql}
            for gq in golden_queries
            if gq.question and gq.generatedSql
        ]
        logger.info("Using %d golden queries as few-shot examples", len(golden_queries_for_prompt))

        # Generate SQL
        try:
            result = _generate_sql(
                question=question,
                entities=entities,
                relationships=relationships,
                glossary_terms=effective_glossary or [],
                dialect=dialect,
                llm_id=llm_id,
                corrections=effective_corrections,
                vocabulary_terms=effective_vocabulary or [],
                golden_queries=golden_queries_for_prompt,
                model_instructions=spec_model.instructions,
                trace_span=span,
            )
            # Include glossary terms in result for extraction/display
            result["glossary_terms_used"] = effective_glossary or []
            result["vocabulary_terms_used"] = effective_vocabulary or []
            _finish_tool_trace(span, output=result)
            return json.dumps(result, indent=2)
        except Exception as e:
            _finish_tool_trace(span, error=str(e))
            raise

    @tool
    def validate_sql(sql: str, callbacks: Callbacks = None) -> str:
        """Validate SQL syntax and security.

        Args:
            sql: The SQL query to validate

        Checks that the SQL is syntactically valid and only reads data
        (no INSERT/UPDATE/DELETE). Always call this before execute_sql.
        """
        from semantic_models.agent.tools import validate_sql as _validate_sql

        span = _start_tool_trace("validate_sql", {"sql": sql, "dialect": dialect}, callbacks)
        try:
            result = _validate_sql(sql, dialect, spec_model)
            _finish_tool_trace(span, output=result)
            return json.dumps(result, indent=2)
        except Exception as e:
            _finish_tool_trace(span, error=str(e))
            raise

    @tool
    def execute_sql(sql: str, callbacks: Callbacks = None) -> str:
        """Execute a SQL query and return results.

        Args:
            sql: The validated SQL query to execute

        Returns columns, rows, row count, and whether results were truncated.
        Only call this after validate_sql returns valid=true.
        """
        from semantic_models.agent.tools import execute_sql as _execute_sql

        span = _start_tool_trace(
            "execute_sql",
            {"sql": sql, "max_rows": ctx._resolved_max_rows},
            callbacks,
        )
        try:
            result = _execute_sql(sql, spec_model, ctx._resolved_max_rows)
            _finish_tool_trace(span, output=result)
            return json.dumps(result, indent=2)
        except Exception as e:
            _finish_tool_trace(span, error=str(e))
            raise

    return [
        list_entities,
        get_entity_attributes,
        get_attribute_values,
        get_relationships,
        resolve_values,
        get_glossary_terms,
        generate_sql,
        validate_sql,
        execute_sql,
    ]


def make_call_model(tools: list):
    """Create the LLM node bound to tools."""
    from dataiku.langchain.dku_llm import DKUChatModel

    def _call_model(state: AgentState, config: RunnableConfig | None = None) -> dict:
        logger.info("-" * 60)
        logger.info("AGENT NODE: DKU_AGENT_CALL_LLM")
        logger.info("  Message count: %d", len(state["messages"]))

        llm = DKUChatModel(llm_id=state["llm_id"])
        model_with_tools = llm.bind_tools(tools)

        response = model_with_tools.invoke(state["messages"], config=config)

        logger.info("  Response type: %s", type(response).__name__)
        if hasattr(response, "tool_calls") and response.tool_calls:
            logger.info("  Tool calls: %s", [tc["name"] for tc in response.tool_calls])
            try:
                tool_calls_text = json.dumps(response.tool_calls, indent=2, ensure_ascii=False, default=str)
            except Exception:
                tool_calls_text = str(response.tool_calls)
            logger.info("  Tool call payload:\n%s", tool_calls_text)
        else:
            logger.info("  No tool calls (final response)")
            logger.info("  Final response content:\n%s", getattr(response, "content", str(response)))

        return {"messages": [response]}

    return _call_model


def make_tools_node(tools: list, ctx: RuntimeContext):
    """Create the ToolNode wrapped with a DKU_AGENT_TOOL_CALLS span."""
    tool_node = ToolNode(tools)

    def _run_tools(state: AgentState, config: RunnableConfig | None = None) -> dict:
        logger.info("-" * 60)
        logger.info("AGENT NODE: DKU_AGENT_TOOL_CALLS")

        callbacks = None
        if config:
            callbacks = config.get("callbacks")

        parent = None
        if callbacks:
            try:
                from dataiku.langchain import dku_span_builder_for_callbacks

                parent = dku_span_builder_for_callbacks(callbacks, ignore_missing=True)
            except Exception:
                parent = None

        tool_calls = []
        last_message = state["messages"][-1]
        if hasattr(last_message, "tool_calls") and last_message.tool_calls:
            tool_calls = last_message.tool_calls

        tool_calls_span = None
        needs_close = False
        if parent is not None:
            tool_calls_span = parent
        elif ctx.trace_span is not None:
            tool_calls_span = ctx.trace_span.subspan("DKU_AGENT_TOOL_CALLS")
            tool_calls_span.__enter__()
            needs_close = True

        if tool_calls_span is not None:
            tool_calls_span.attributes["nbToolCalls"] = len(tool_calls)
            ctx.tool_calls_span = tool_calls_span

        try:
            return tool_node.invoke(state, config=config)
        finally:
            if needs_close and tool_calls_span is not None:
                tool_calls_span.__exit__(None, None, None)
            ctx.tool_calls_span = None

    return _run_tools


def should_continue(state: AgentState) -> Literal["tools", "end"]:
    """Determine if we should continue with tools or end."""
    last_message = state["messages"][-1]

    # If the model made tool calls, continue
    if hasattr(last_message, "tool_calls") and last_message.tool_calls:
        return "tools"

    # Otherwise, end
    return "end"


def build_agent_graph(call_model_node, tools_node) -> StateGraph:
    """Build the LangGraph agent."""
    graph = StateGraph(AgentState)

    # Add nodes
    graph.add_node("DKU_AGENT_CALL_LLM", call_model_node)
    graph.add_node("DKU_AGENT_TOOL_CALLS", tools_node)

    # Set entry point
    graph.set_entry_point("DKU_AGENT_CALL_LLM")

    # Add conditional edges
    graph.add_conditional_edges(
        "DKU_AGENT_CALL_LLM",
        should_continue,
        {
            "tools": "DKU_AGENT_TOOL_CALLS",
            "end": END,
        },
    )

    # Tools always go back to agent
    graph.add_edge("DKU_AGENT_TOOL_CALLS", "DKU_AGENT_CALL_LLM")

    return graph.compile()


def run_agent_core(ctx: RuntimeContext) -> dict[str, Any]:
    from semantic_models.agent.prompts.system import get_system_prompt

    start_time = time.time()

    # Initialize trace span (use provided root if available)
    if ctx.trace is not None and hasattr(ctx.trace, "subspan"):
        ctx.trace_root = ctx.trace
        ctx.trace_span = ctx.trace.subspan("semantic-model-query")
    else:
        if ctx.trace is not None and not hasattr(ctx.trace, "subspan"):
            logger.warning("Provided trace does not support subspan; creating new trace")
        ctx.trace_root = new_trace("semantic-model-query")
        ctx.trace_span = ctx.trace_root

    logger.info("=" * 80)
    logger.info("AGENT START")
    logger.info("=" * 80)
    logger.info("Question: %s", ctx.question)
    logger.info("Semantic Model: %s", ctx.semantic_model_id)
    logger.info("Dialect: %s", ctx.dialect)
    logger.info("LLM: %s", ctx.llm_id)

    final_state: AgentState | None = None
    result: dict[str, Any] | None = None
    total_duration_ms = 0.0

    with ctx.trace_span:
        ctx.trace_span.inputs["question"] = ctx.question
        ctx.trace_span.attributes["semantic_model_id"] = ctx.semantic_model_id
        if ctx.version_id is not None:
            ctx.trace_span.attributes["version_id"] = ctx.version_id
        ctx.trace_span.attributes["dialect"] = ctx.dialect
        ctx.trace_span.attributes["llm_id"] = ctx.llm_id
        if ctx.embedding_llm_id is not None:
            ctx.trace_span.attributes["embedding_llm_id"] = ctx.embedding_llm_id

        try:
            # Build tools and graph
            tools = create_tools(ctx)
            agent = build_agent_graph(make_call_model(tools), make_tools_node(tools, ctx))

            # Prepare system prompt
            system_prompt = get_system_prompt(ctx.dialect)
            logger.info("System prompt:\n%s", system_prompt)

            initial_state: AgentState = {
                "question": ctx.question,
                "semantic_model_id": ctx.semantic_model_id,
                "version_id": ctx.version_id,
                "spec": ctx.spec,
                "dialect": ctx.dialect,
                "llm_id": ctx.llm_id,
                "ctx": ctx,
                "messages": [
                    SystemMessage(content=system_prompt),
                    HumanMessage(content=ctx.question),
                ],
                "sql_attempts": [],
                "sql_execution_count": 0,
                "start_time": start_time,
                "final_result": None,
            }

            # Set up tracing callback
            callbacks = []
            if ctx.trace_span is not None:
                try:
                    from dataiku.langchain import LangchainToDKUTracer

                    tracer = LangchainToDKUTracer(dku_trace=ctx.trace_span)
                    callbacks.append(tracer)
                    logger.info("DSS tracing enabled")
                except ImportError:
                    logger.warning("LangchainToDKUTracer not available, tracing disabled")

            # Run the agent with recursion limit
            recursion_limit = ctx._resolved_recursion_limit

            logger.info("-" * 60)
            logger.info("INVOKING AGENT GRAPH (recursion_limit=%d)", recursion_limit)
            logger.info("-" * 60)

            invoke_config = {"recursion_limit": recursion_limit}
            if callbacks:
                invoke_config["callbacks"] = callbacks

            final_state = agent.invoke(initial_state, config=invoke_config)

            # Extract results and build sources/artifacts
            total_duration_ms = (time.time() - start_time) * 1000
            result = _extract_results(final_state, ctx, total_duration_ms)

            if result["success"]:
                logger.info("=" * 80)
                logger.info("AGENT SUCCESS - %.0fms", total_duration_ms)
                logger.info("  Rows returned: %d", result.get("row_count", 0))
                logger.info("  SQL executions: %d", len(result.get("sql_executions", [])))
                logger.info("=" * 80)
            else:
                logger.info("=" * 80)
                logger.info("AGENT COMPLETED - No SQL results - %.0fms", total_duration_ms)
                logger.info("=" * 80)

            if ctx.trace_span is not None:
                ctx.trace_span.outputs["success"] = result.get("success")
                if result.get("sql") is not None:
                    ctx.trace_span.outputs["sql"] = result.get("sql")
                if result.get("row_count") is not None:
                    ctx.trace_span.outputs["row_count"] = result.get("row_count")

        except Exception as e:
            total_duration_ms = (time.time() - start_time) * 1000
            logger.exception("AGENT EXCEPTION")
            ctx.add_error(str(e))
            if ctx.trace_span is not None:
                ctx.trace_span.attributes["status"] = "error"
                ctx.trace_span.outputs["error"] = str(e)
            result = {
                "success": False,
                "error": str(e),
                "sources": ctx.sources,
                "artifacts": [],
            }

    if result is None:
        result = {
            "success": False,
            "error": "Agent failed unexpectedly",
            "sources": ctx.sources,
            "artifacts": [],
        }

    message_count = None
    if isinstance(result.get("trace"), dict):
        message_count = result["trace"].get("message_count")
    if message_count is None and final_state is not None:
        message_count = len(final_state["messages"])

    result["trace"] = {
        "message_count": message_count,
        "total_duration_ms": total_duration_ms,
        "span": ctx.trace_span.to_dict() if ctx.trace_span is not None else None,
    }
    return result


def _extract_results(
    final_state: AgentState,
    ctx: RuntimeContext,
    total_duration_ms: float,
) -> dict[str, Any]:
    messages = final_state["messages"]

    # Get the final AI message (agent's natural language response)
    final_message = messages[-1]
    final_response = ""
    if isinstance(final_message, AIMessage):
        final_response = final_message.content if hasattr(final_message, "content") else str(final_message)

    # Collect resolve_values corrections and glossary terms used
    corrections = []
    glossary_terms_used = []
    glossary_keys: set[str] = set()
    glossary_from_tool = False

    def _glossary_key(term: dict) -> str | None:
        term_id = term.get("term_id") or term.get("id")
        if term_id:
            return f"id:{term_id}"
        term_value = term.get("term")
        if isinstance(term_value, str) and term_value.strip():
            return f"term:{term_value.strip().lower()}"
        return None

    for msg in messages:
        if isinstance(msg, ToolMessage):
            try:
                result = json.loads(msg.content)

                # Capture resolve_values corrections
                if msg.name == "resolve_values" and result.get("corrections"):
                    for c in result["corrections"]:
                        # Avoid duplicates
                        if not any(
                            existing["original"] == c.get("original") and existing["resolved"] == c.get("resolved")
                            for existing in corrections
                        ):
                            corrections.append(
                                {
                                    "original": c.get("original"),
                                    "resolved": c.get("resolved"),
                                    "entity_id": c.get("entity_id"),
                                    "attribute": c.get("attribute"),
                                    "score": c.get("score"),
                                    "match_type": c.get("match_type"),
                                    "llm_inferred": c.get("llm_inferred", False),
                                }
                            )

                # Capture get_glossary_terms results (already filtered at tool level)
                elif msg.name == "get_glossary_terms" and isinstance(result, list):
                    for term in result:
                        key = _glossary_key(term)
                        if key and key in glossary_keys:
                            continue
                        if key:
                            glossary_keys.add(key)
                        glossary_terms_used.append(
                            {
                                "term_id": term.get("term_id"),
                                "term": term.get("term"),
                                "description": term.get("description"),
                                "synonyms": term.get("synonyms"),
                                "entity_id": term.get("entity_id"),
                                "entity_name": term.get("entity_name"),
                                "target_type": term.get("target_type"),
                                "target_id": term.get("target_id"),
                                "target_name": term.get("target_name"),
                                "is_primary": term.get("is_primary"),
                            }
                        )
                    if result:
                        glossary_from_tool = True

                # Also capture glossary terms from generate_sql (auto-loaded)
                elif (
                    msg.name == "generate_sql"
                    and result.get("glossary_terms_used")
                    and not glossary_from_tool
                ):
                    for term in result["glossary_terms_used"]:
                        key = _glossary_key(term)
                        if key and key in glossary_keys:
                            continue
                        if key:
                            glossary_keys.add(key)
                        glossary_terms_used.append(
                            {
                                "term_id": term.get("term_id"),
                                "term": term.get("term"),
                                "description": term.get("description"),
                                "synonyms": term.get("synonyms"),
                                "entity_id": term.get("entity_id"),
                                "entity_name": term.get("entity_name"),
                                "target_type": term.get("target_type"),
                                "target_id": term.get("target_id"),
                                "target_name": term.get("target_name"),
                                "is_primary": term.get("is_primary"),
                            }
                        )

            except json.JSONDecodeError:
                pass

    logger.info(
        "Extracted %d corrections, %d glossary terms matching question", len(corrections), len(glossary_terms_used)
    )

    # Collect ALL execute_sql calls and their results (for debugging/multi-SQL)
    sql_executions = []
    for i, msg in enumerate(messages):
        if isinstance(msg, ToolMessage) and msg.name == "execute_sql":
            try:
                result = json.loads(msg.content)

                # Find the SQL from the preceding tool call by matching tool_call_id
                tool_call_id = getattr(msg, "tool_call_id", None)
                sql_query = None
                for prev_msg in reversed(messages[:i]):
                    if isinstance(prev_msg, AIMessage) and hasattr(prev_msg, "tool_calls"):
                        for tc in prev_msg.tool_calls or []:
                            # Match by tool_call_id if available, otherwise fall back to name match
                            if tool_call_id and tc.get("id") == tool_call_id:
                                sql_query = tc["args"].get("sql")
                                break
                            elif not tool_call_id and tc["name"] == "execute_sql":
                                # Fallback for older messages without tool_call_id
                                sql_query = tc["args"].get("sql")
                                break
                        if sql_query:
                            break

                sql_executions.append(
                    {
                        "sql": sql_query,
                        "success": result.get("success", False),
                        "columns": result.get("columns", []),
                        "rows": result.get("rows", []),
                        "row_count": result.get("row_count", 0),
                        "truncated": result.get("truncated", False),
                        "error": result.get("error"),
                    }
                )

            except json.JSONDecodeError:
                pass

    # Find the last successful execution for primary result
    last_success = None
    for exec_item in reversed(sql_executions):
        if exec_item["success"]:
            last_success = exec_item
            break

    # Build sources for ALL successful executions (for true multi-SQL support)
    successful_executions = [e for e in sql_executions if e["success"]]
    for exec_item in successful_executions:
        ctx.add_sql_query(exec_item["sql"])
        ctx.add_records(exec_item["columns"], exec_item["rows"])

    # Build artifacts for ALL successful results (using context's limits)
    artifacts = []
    if ctx._resolved_artifacts_limit > 0:
        for i, exec_item in enumerate(successful_executions):
            # For single SQL, use simple name; for multi-SQL, number them
            if len(successful_executions) == 1:
                artifact_name = "Query Results"
            else:
                artifact_name = f"Query Results ({i + 1})"

            artifact = ctx.build_records_artifact(
                name=artifact_name,
                description="",
                columns=exec_item["columns"],
                rows=exec_item["rows"],
            )
            if artifact:
                artifacts.append(artifact)

    # Build response - use last_success for primary result (backward compat)
    if successful_executions:
        return {
            "success": True,
            "response": final_response,
            "sql": last_success["sql"],
            "columns": last_success["columns"],
            "rows": last_success["rows"],
            "row_count": last_success["row_count"],
            "truncated": last_success["truncated"],
            "sql_executions": sql_executions,
            "corrections": corrections,
            "glossary_terms_used": glossary_terms_used,
            "sources": ctx.sources,
            "artifacts": artifacts,
            "trace": {
                "message_count": len(messages),
                "total_duration_ms": total_duration_ms,
            },
        }
    else:
        # No successful SQL execution
        has_failed_sql = any(not e["success"] for e in sql_executions)

        if has_failed_sql or not final_response:
            # Actual error - SQL execution failed or no response at all
            error_msg = final_response or "Agent completed without successful SQL execution"
            ctx.add_error(error_msg)
        else:
            # Agent provided an answer without needing SQL (informational response)
            ctx.add_info(f"Agent response (no SQL needed): {final_response}")

        # Try to get partial SQL
        partial_sql = None
        for exec_item in reversed(sql_executions):
            if exec_item["sql"]:
                partial_sql = exec_item["sql"]
                break

        return {
            "success": False,
            "error": final_response if (has_failed_sql or not final_response) else None,
            "response": final_response,
            "sql": partial_sql,
            "sql_executions": sql_executions,
            "corrections": corrections,
            "glossary_terms_used": glossary_terms_used,
            "sources": ctx.sources,
            "artifacts": [],
            "trace": {
                "message_count": len(messages),
                "total_duration_ms": total_duration_ms,
            },
        }


# =============================================================================
# Public Entry Point
# =============================================================================


def run_agent(
    question: str,
    semantic_model_id: str,
    spec: dict,
    dialect: str,
    llm_id: str,
    embedding_llm_id: str | None = None,
    version_id: str | None = None,
    trace: Any | None = None,
) -> dict[str, Any]:
    ctx = RuntimeContext(
        question=question,
        spec=spec,
        dialect=dialect,
        llm_id=llm_id,
        embedding_llm_id=embedding_llm_id,
        semantic_model_id=semantic_model_id,
        version_id=version_id,
        trace=trace,
    )
    return run_agent_core(ctx)
