import logging
import time
import uuid
from typing import Any, Dict

from dataiku.llm.agent_tools import BaseAgentTool
from semantic_models.agent.semantic_model_query import (
    build_client_error_response,
    format_tool_response,
    run_semantic_model_query,
)

logger = logging.getLogger(__name__)


class SemanticModelQueryTool(BaseAgentTool):
    def set_config(self, config: Dict[str, Any], plugin_config: Dict[str, Any]) -> None:
        from semantic_models.logging_utils import configure_logging

        configure_logging()

        config = dict(config)
        self.config = config
        self.plugin_config = plugin_config

        # Project key (optional - defaults to current project)
        self.project_key = config.get("project_key")

        self.semantic_model_id = config.get("semantic_model_id")
        if not self.semantic_model_id:
            raise ValueError("semantic_model_id is required")

        version_str = config.get("version_number", "")
        self.version_id = version_str if version_str else None

        self.llm_id = config.get("llm_id")
        if not self.llm_id:
            raise ValueError("llm_id is required")

        from semantic_models.services.factory import get_semantic_models_service
        from semantic_models.services.local_semantic_model_client import LocalSemanticModelClient

        client = LocalSemanticModelClient(self.project_key)
        self.service = get_semantic_models_service(client=client)

        sql_generation_mode = config.get("sql_generation_mode", "VERSION")

        # Load version-level SQL generation config (if available)
        try:
            if self.version_id is not None:
                version = self.service.get_version(self.semantic_model_id, self.version_id)
            else:
                version = self.service.get_active_version(self.semantic_model_id)
            version_cfg = getattr(version, "sqlGenerationConfig", {}) or {}
        except Exception as e:
            logger.warning("Failed to load version for SQL generation defaults: %s", e)
            version_cfg = {}

        defaults = {
            "max_rows_per_query": 1000,
            "agent_recursion_limit": 100,
            "sources_records_limit": 50,
            "artifacts_records_limit": 1000,
        }

        if sql_generation_mode == "CUSTOM":
            logger.info("SQL generation parameters source: tool configuration")
            for key, default in defaults.items():
                if config.get(key) is None:
                    config[key] = default
                    logger.info("SQL param %s=%s (source=default)", key, config[key])
                else:
                    logger.info("SQL param %s=%s (source=tool)", key, config[key])
        else:
            logger.info("SQL generation parameters source: semantic model version")
            for key, default in defaults.items():
                if key in version_cfg:
                    config[key] = version_cfg.get(key)
                    logger.info("SQL param %s=%s (source=version)", key, config[key])
                else:
                    config[key] = default
                    logger.info("SQL param %s=%s (source=default)", key, config[key])

        raw_embedding = config.get("embedding_llm_id")
        if isinstance(raw_embedding, str) and raw_embedding.strip():
            self.embedding_llm_id = raw_embedding.strip()
            logger.info("Embedding LLM source: tool configuration (%s)", self.embedding_llm_id)
        else:
            self.embedding_llm_id = None
            logger.info("Embedding LLM source: none (semantic search disabled)")

        from semantic_models.config import load_tool_config

        load_tool_config(config)

        logger.info(
            "SemanticModelQueryTool configured: project=%s, model=%s, llm=%s, embedding_llm=%s",
            self.project_key or "(current)",
            self.semantic_model_id,
            self.llm_id,
            self.embedding_llm_id,
        )

        # Initialize search indices (value + semantic) for the configured model/version
        self._initialize_search_indices()

    def _initialize_search_indices(self) -> None:
        """Initialize value and semantic indices for the configured model."""
        try:
            if self.version_id is not None:
                version = self.service.get_version(self.semantic_model_id, self.version_id)
            else:
                version = self.service.get_active_version(self.semantic_model_id)

            if not version:
                return

            from semantic_models.resolution.index import (
                collect_sources_from_spec,
                initialize_value_indices_from_sources,
            )
            from semantic_models.utils.spec_utils import build_spec_from_version

            spec_dict = build_spec_from_version(version).model_dump()
            sources = collect_sources_from_spec(self.semantic_model_id, version.id, spec_dict)

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

        except Exception as e:
            logger.warning("Failed to initialize search indices: %s", e)
            # Continue anyway - exact/fuzzy matching still works

    def get_descriptor(self, tool: Any) -> Dict[str, Any]:
        # Fetch model info for better description
        try:
            model_with_versions = self.service.get_model(self.semantic_model_id)
            model = model_with_versions.model
            model_name = model.name
            model_desc = model.description or ""
        except Exception:
            model_name = self.semantic_model_id
            model_desc = ""

        description = f"""Query the '{model_name}' semantic model using natural language.

{model_desc}

This tool will:
1. Analyze your question to find relevant data entities
2. Generate appropriate SQL query
3. Execute the query against the data source
4. Return results and/or a natural language answer

Use this for data questions about the semantic model."""

        return {
            "description": description,
            "inputSchema": {
                "$id": "https://dataiku.com/agents/tools/semantic-model-query/input",
                "title": "Semantic Model Query Input",
                "type": "object",
                "properties": {"question": {"type": "string", "description": "Natural language question to answer"}},
                "required": ["question"],
            },
        }

    def invoke(self, input: Dict[str, Any], trace: Any) -> Dict[str, Any]:
        question = input.get("input", {}).get("question")
        if not question:
            raise ValueError("Question is required")

        from semantic_models.logging_utils import log_context

        request_id = str(uuid.uuid4())
        started_at = time.perf_counter()
        user = self._resolve_user()
        with log_context(request_id=request_id, user=user, method="TOOL", path="/agent-tools/semantic-model-query"):
            logger.info("SemanticModelQueryTool invoked: %s", question[:100])
            logger.info(
                "Invocation context: project=%s model=%s version=%s llm=%s embedding_llm=%s question_chars=%d",
                self.project_key or "(current)",
                self.semantic_model_id,
                self.version_id or "(active)",
                self.llm_id,
                self.embedding_llm_id or "(none)",
                len(question),
            )

            try:
                result = run_semantic_model_query(
                    service=self.service,
                    question=question,
                    semantic_model_id=self.semantic_model_id,
                    version_id=self.version_id,
                    llm_id=self.llm_id,
                    embedding_llm_id=self.embedding_llm_id,
                    initialize_indices=False,
                    trace=trace,
                )
                duration_ms = (time.perf_counter() - started_at) * 1000
                logger.info(
                    "SemanticModelQueryTool completed: status=success duration_ms=%.2f success=%s row_count=%s",
                    duration_ms,
                    result.get("success"),
                    result.get("row_count"),
                )

                return format_tool_response(
                    question=question,
                    result=result,
                    wrap_sources=True,
                )

            except Exception as e:
                duration_ms = (time.perf_counter() - started_at) * 1000
                logger.exception("SemanticModelQueryTool completed: status=error duration_ms=%.2f", duration_ms)
                return build_client_error_response(question, str(e), wrap_sources=True)

    def load_sample_query(self, tool) -> Dict[str, Any]:
        return {"input": {"question": "What are the top 10 records?"}}

    @staticmethod
    def _resolve_user() -> str:
        """Best effort DSS user resolution for tool invocations."""
        try:
            import dataiku

            auth = dataiku.api_client().get_auth_info()
            return auth.get("authIdentifier") or "-"
        except Exception:
            return "-"
