import logging
from typing import Any, Optional

from semantic_models.agent.autogen_tools import AutogenTools, ToolResult
from semantic_models.services.dataset_metadata import get_dataset_profile as _get_dataset_profile
from semantic_models.services.dss_utils import fetch_distinct_values
from semantic_models.utils import quote_identifier

logger = logging.getLogger(__name__)

MAX_SAMPLE_VALUE_LENGTH = 100
MAX_SAMPLE_VALUES = 20


def _truncate_value(value: Any, max_len: int = MAX_SAMPLE_VALUE_LENGTH) -> str:
    text = str(value)
    if len(text) > max_len:
        return text[: max_len - 3] + "..."
    return text


def _truncate_values(values: list[Any], max_values: int = MAX_SAMPLE_VALUES) -> list[str]:
    trimmed = values[:max_values]
    return [_truncate_value(v) for v in trimmed]


class AgenticTools:
    def __init__(self, ctx: Any):
        self.ctx = ctx
        self.autogen = AutogenTools(ctx)

    def get_dataset_profile(
        self,
        project_key: str,
        dataset_name: str,
        sample_value_limit: int = 0,
    ) -> ToolResult:
        try:
            profile = _get_dataset_profile(project_key, dataset_name, sample_value_limit=sample_value_limit)
            # Populate schema cache for downstream tools
            self.ctx.schema_cache = {
                "columns": [
                    {
                        "name": c["name"],
                        "type": c.get("type", "string"),
                        "description": c.get("description"),
                    }
                    for c in profile.get("columns", [])
                ],
                "connection": profile.get("connection"),
                "full_table": profile.get("full_table"),
                "dialect": profile.get("dialect", "postgres"),
            }
            # Seed profile cache with any sample values (if provided)
            for c in profile.get("columns", []):
                col_name = c.get("name")
                if not col_name:
                    continue
                sample_values = c.get("sample_values")
                if sample_values:
                    self.ctx.profile_cache[col_name] = {
                        "sample_values": sample_values,
                        "is_categorical": True,
                    }
            return ToolResult(success=True, data=profile)
        except Exception as e:
            logger.exception("Failed to get dataset profile")
            return ToolResult(success=False, error=str(e))

    def get_column_samples(self, column_name: str, limit: int = 10) -> ToolResult:
        if not self.ctx.schema_cache:
            return ToolResult(success=False, error="Schema not loaded")
        try:
            values = fetch_distinct_values(
                connection=self.ctx.schema_cache["connection"],
                full_table=self.ctx.schema_cache["full_table"],
                column_name=column_name,
                dialect=self.ctx.schema_cache.get("dialect", "postgres"),
                limit=limit,
                max_value_length=MAX_SAMPLE_VALUE_LENGTH,
            )
            values = _truncate_values(values)
            # Update profile cache with sample values
            self.ctx.profile_cache[column_name] = {
                **self.ctx.profile_cache.get(column_name, {}),
                "sample_values": values,
                "is_categorical": bool(values),
            }
            return ToolResult(success=True, data={"column": column_name, "values": values})
        except Exception as e:
            logger.warning("Failed to fetch samples for %s: %s", column_name, e)
            return ToolResult(success=False, error=str(e))

    def get_column_stats(self, column_name: str) -> ToolResult:
        if not self.ctx.schema_cache:
            return ToolResult(success=False, error="Schema not loaded")
        try:
            import dataiku

            connection = self.ctx.schema_cache["connection"]
            full_table = self.ctx.schema_cache["full_table"]
            dialect = self.ctx.schema_cache.get("dialect", "postgres")

            col_sql = quote_identifier(column_name, dialect)
            sql = (
                f"SELECT COUNT(*) AS total, "
                f"COUNT({col_sql}) AS non_null, "
                f"COUNT(DISTINCT {col_sql}) AS distinct_count "
                f"FROM {full_table}"
            )
            executor = dataiku.core.sql.SQLExecutor2(connection=connection)
            df = executor.query_to_df(sql)
            if df.empty:
                stats = {"null_count": 0, "distinct_count": 0}
            else:
                total = int(df.iloc[0]["total"])
                non_null = int(df.iloc[0]["non_null"])
                distinct_count = int(df.iloc[0]["distinct_count"])
                stats = {
                    "null_count": max(0, total - non_null),
                    "distinct_count": distinct_count,
                }

            self.ctx.profile_cache[column_name] = {
                **self.ctx.profile_cache.get(column_name, {}),
                **stats,
            }
            return ToolResult(success=True, data={"column": column_name, **stats})
        except Exception as e:
            logger.warning("Failed to fetch stats for %s: %s", column_name, e)
            return ToolResult(success=False, error=str(e))

    # Delegate model mutation/validation to AutogenTools
    def create_entity(self, entity_spec: dict[str, Any]) -> ToolResult:
        return self.autogen.create_entity(entity_spec)

    def add_metrics(self, entity_id: str, metrics: list[dict[str, Any]]) -> ToolResult:
        return self.autogen.add_metrics(entity_id, metrics)

    def add_filters(self, entity_id: str, filters: list[dict[str, Any]]) -> ToolResult:
        return self.autogen.add_filters(entity_id, filters)

    def add_relationship(
        self,
        left_entity: str,
        right_entity: str,
        left_column: str,
        right_column: str,
    ) -> ToolResult:
        return self.autogen.add_relationship(left_entity, right_entity, left_column, right_column)

    def coverage_report(self) -> ToolResult:
        return self.autogen.coverage_report()

    def validate_model(self) -> ToolResult:
        return self.autogen.validate_model()
