import logging
from typing import Any, Optional

from pydantic import BaseModel

from semantic_models.agent.generation_state import AutogenContext

logger = logging.getLogger(__name__)


class ToolResult(BaseModel):
    """Standard tool result format."""

    success: bool
    data: Any = None
    error: Optional[str] = None


class AutogenTools:
    """
    Tool implementations for autogen agent.

    Each tool operates on the autogen context and returns a ToolResult.
    The runner calls these tools during pipeline execution.
    """

    def __init__(self, context: AutogenContext):
        """
        Initialize tools with autogen context.

        Args:
            context: The context to operate on
        """
        self.ctx = context

    # =========================================================================
    # Discovery & Profiling Tools
    # =========================================================================

    def get_dataset_schema(self) -> ToolResult:
        """
        Get columns, types, and descriptions for the dataset.

        Returns cached schema if available, otherwise fetches from DSS.
        """
        from semantic_models.services.dataset_profiling import get_dataset_schema

        if self.ctx.schema_cache:
            return ToolResult(success=True, data=self.ctx.schema_cache)

        try:
            schema = get_dataset_schema(self.ctx.project_key, self.ctx.dataset_name)
            self.ctx.schema_cache = schema
            return ToolResult(success=True, data=schema)
        except Exception as e:
            logger.exception("Failed to get dataset schema")
            return ToolResult(success=False, error=str(e))

    def get_column_profile(self, column_name: str) -> ToolResult:
        """
        Get statistics for a column from the profile cache.

        Profile data is populated by the profile_columns step during batch profiling.
        """
        if column_name in self.ctx.profile_cache:
            return ToolResult(success=True, data=self.ctx.profile_cache[column_name])

        return ToolResult(
            success=False,
            error=f"Column '{column_name}' not found in profile cache. Run profile_columns first.",
        )

    def get_top_values(self, column_name: str, limit: int = 5) -> ToolResult:
        """
        Get sample distinct values for a column.

        Queries the database directly - use sparingly for columns needing more
        sample values than captured during batch profiling.
        """
        from semantic_models.services.dataset_profiling import get_top_values as _get_top_values

        if not self.ctx.schema_cache:
            return ToolResult(success=False, error="Schema not loaded")

        try:
            values = _get_top_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,
            )
            return ToolResult(success=True, data=values)
        except Exception as e:
            logger.warning(f"Failed to get top values for {column_name}: {e}")
            return ToolResult(success=False, error=str(e))

    # =========================================================================
    # Modeling Tools
    # =========================================================================

    def create_entity(self, entity_spec: dict[str, Any]) -> ToolResult:
        """
        Create or update (upsert) an entity in the working model.

        Args:
        entity_spec: Dict with {entity_id, entity_name, description, columns, primaryKey}

        Returns:
            ToolResult with the created entity ID
        """
        from semantic_models.schemas import Attribute, EntityClass

        entity_id = entity_spec.get("entity_id")
        if not entity_id:
            return ToolResult(success=False, error="entity_id is required")

        # Build attributes from columns, enriched with profile data
        attributes = []
        for col_name in entity_spec.get("columns", []):
            # Get type from schema cache
            col_type = "string"
            col_description = None
            if self.ctx.schema_cache:
                for col in self.ctx.schema_cache.get("columns", []):
                    if col["name"] == col_name:
                        col_type = col["type"]
                        raw_description = col.get("description")
                        if isinstance(raw_description, str) and raw_description.strip():
                            col_description = raw_description.strip()
                        break

            # Determine resolvable from profile cache
            resolvable = False
            if col_name in self.ctx.profile_cache:
                profile = self.ctx.profile_cache[col_name]
                # Generated attributes should default to AUTO_INDEX for categorical columns.
                resolvable = bool(profile.get("is_categorical", False))

            attributes.append(
                Attribute(
                    name=col_name,
                    description=col_description,
                    dssType=col_type,
                    type="COLUMN",
                    column=col_name,
                    manualValues=[],
                    resolveInUserRequests=resolvable,
                    distinctValuesHandlingMode="AUTO_INDEX" if resolvable else "NONE",
                    indexDistinctValues=resolvable,
                )
            )

        # Use LLM-provided primaryKey only (no heuristic fallback)
        primary_key_cols = entity_spec.get("primaryKey", [])

        display_name = entity_spec.get("entity_name")
        description = entity_spec.get("description") or (
            display_name if display_name and display_name != entity_id else None
        )

        entity = EntityClass(
            name=entity_id,
            description=description,
            type="DATASET",
            datasetRef=f"{self.ctx.project_key}.{self.ctx.dataset_name}",
            attributes=attributes,
            metrics=[],
            filters=[],
            primaryKey={"attributes": primary_key_cols},
        )

        # Upsert: remove existing entity with same ID if present
        if not self.ctx.working_model:
            self.ctx.working_model = {"entities": [], "relationships": []}

        self.ctx.working_model["entities"] = [
            e for e in self.ctx.working_model["entities"] if e.get("name") != entity.name
        ]
        self.ctx.working_model["entities"].append(entity.model_dump())

        return ToolResult(success=True, data={"entity_id": entity.name})

    def update_entity(self, entity_id: str, changes: dict[str, Any]) -> ToolResult:
        """
        Update entity properties (name, description, primaryKey).

        Args:
            entity_id: ID of entity to update
            changes: Dict of properties to update

        Returns:
            ToolResult with updated entity
        """
        if not self.ctx.working_model:
            return ToolResult(success=False, error="No working model")

        for entity in self.ctx.working_model["entities"]:
            if entity.get("name") == entity_id:
                if "primaryKey" in changes:
                    if isinstance(changes["primaryKey"], list):
                        entity["primaryKey"] = {"attributes": changes.pop("primaryKey")}
                    else:
                        entity["primaryKey"] = changes.pop("primaryKey")
                entity.update(changes)
                return ToolResult(success=True, data=entity)

        return ToolResult(success=False, error=f"Entity {entity_id} not found")

    def add_metrics(self, entity_id: str, metrics: list[dict[str, Any]]) -> ToolResult:
        """
        Add metrics to an entity.

        Validates that each metric contains an aggregate function.

        Args:
            entity_id: ID of entity to add metrics to
            metrics: List of metric dicts with {name, expression, description}

        Returns:
            ToolResult with count of metrics added
        """
        aggregate_functions = ["SUM", "COUNT", "AVG", "MIN", "MAX"]

        valid_metrics = []
        for m in metrics:
            expression = m.get("pseudoSQLExpression") or m.get("expression", "")
            expr_upper = expression.upper()
            has_agg = any(f"{agg}(" in expr_upper for agg in aggregate_functions)
            if has_agg:
                metric = {
                    "name": m.get("name"),
                    "description": m.get("description"),
                    "pseudoSQLExpression": expression,
                }
                valid_metrics.append(metric)
            else:
                logger.warning(f"Skipping metric without aggregate: {m.get('name')}")

        if not self.ctx.working_model:
            return ToolResult(success=False, error="No working model")

        for entity in self.ctx.working_model.get("entities", []):
            if entity.get("name") == entity_id:
                if "metrics" not in entity:
                    entity["metrics"] = []
                entity["metrics"].extend(valid_metrics)
                return ToolResult(success=True, data={"added": len(valid_metrics)})

        return ToolResult(success=False, error=f"Entity {entity_id} not found")

    def add_filters(self, entity_id: str, filters: list[dict[str, Any]]) -> ToolResult:
        """
        Add filters to an entity.

        Args:
            entity_id: ID of entity to add filters to
            filters: List of filter dicts with {name, expression, description}

        Returns:
            ToolResult with count of filters added
        """
        if not self.ctx.working_model:
            return ToolResult(success=False, error="No working model")

        normalized_filters = []
        for f in filters:
            normalized_filters.append(
                {
                    "name": f.get("name"),
                    "description": f.get("description"),
                    "pseudoSQLExpression": f.get("pseudoSQLExpression") or f.get("expression", ""),
                }
            )

        for entity in self.ctx.working_model.get("entities", []):
            if entity.get("name") == entity_id:
                if "filters" not in entity:
                    entity["filters"] = []
                entity["filters"].extend(normalized_filters)
                return ToolResult(success=True, data={"added": len(normalized_filters)})

        return ToolResult(success=False, error=f"Entity {entity_id} not found")

    def add_relationship(
        self,
        left_entity: str,
        right_entity: str,
        left_column: str,
        right_column: str,
    ) -> ToolResult:
        """
        Add a relationship directly to the working model.

        Args:
            left_entity: ID of left entity
            right_entity: ID of right entity
            left_column: Column in left entity
            right_column: Column in right entity

        Returns:
            ToolResult with the relationship
        """
        if not self.ctx.working_model:
            self.ctx.working_model = {"entities": [], "relationships": []}

        entity_names = {e.get("name") for e in self.ctx.working_model.get("entities", [])}
        if left_entity not in entity_names or right_entity not in entity_names:
            logger.warning(
                "Skipping relationship with unknown entities: %s <-> %s",
                left_entity,
                right_entity,
            )
            return ToolResult(success=False, error="Relationship references unknown entity")

        relationship = {
            "firstEntity": left_entity,
            "secondEntity": right_entity,
            "pseudoSQLExpression": f"{left_entity}.{left_column} = {right_entity}.{right_column}",
        }

        self.ctx.working_model["relationships"].append(relationship)
        return ToolResult(success=True, data=relationship)

    # =========================================================================
    # Validation Tools
    # =========================================================================

    def coverage_report(self) -> ToolResult:
        """
        Compute column coverage from entity_groups.

        Coverage is computed dynamically from current entity_groups,
        not stored in a coverage_map. This properly detects duplicates
        by tracking all assignments.

        Returns:
            ToolResult with coverage report dict
        """
        if not self.ctx.schema_cache:
            return ToolResult(success=False, error="Schema not loaded")

        all_columns = [c["name"] for c in self.ctx.schema_cache.get("columns", [])]

        # Build column -> [entity_ids] mapping from entity_groups
        column_assignments: dict[str, list[str]] = {}
        for group in self.ctx.entity_groups:
            entity_id = group.get("entity_id", "")
            for col in group.get("columns", []):
                if col not in column_assignments:
                    column_assignments[col] = []
                column_assignments[col].append(entity_id)

        # Compute coverage stats
        all_columns_set = set(all_columns)
        covered_columns = set(column_assignments.keys())
        unassigned = sorted(list(all_columns_set - covered_columns))
        shared = sorted([col for col, entities in column_assignments.items() if len(entities) > 1])
        shared_details = {col: entities for col, entities in column_assignments.items() if len(entities) > 1}

        total = len(all_columns)
        used = len(covered_columns & all_columns_set)

        report = {
            "total_columns": total,
            "used_columns": used,
            "unassigned_columns": unassigned,
            "shared_columns": shared,
            "shared_details": shared_details,
        }

        return ToolResult(success=True, data=report)

    def validate_model(self) -> ToolResult:
        errors: list[str] = []
        warnings: list[str] = []

        if not self.ctx.working_model:
            return ToolResult(success=False, error="No working model")

        # Check coverage - report as warnings, not errors
        coverage_result = self.coverage_report()
        if coverage_result.success and coverage_result.data:
            coverage = coverage_result.data
            if coverage["unassigned_columns"]:
                warnings.append(f"{len(coverage['unassigned_columns'])} columns are unassigned")
            if coverage["shared_columns"]:
                warnings.append(f"{len(coverage['shared_columns'])} columns are shared across entities")

        # Check metrics have aggregates
        aggregate_functions = ["SUM", "COUNT", "AVG", "MIN", "MAX"]
        for entity in self.ctx.working_model.get("entities", []):
            for metric in entity.get("metrics", []):
                expr_upper = (metric.get("pseudoSQLExpression") or "").upper()
                has_agg = any(f"{agg}(" in expr_upper for agg in aggregate_functions)
                if not has_agg:
                    errors.append(f"Metric '{metric.get('name')}' missing aggregate function")

        # Check relationships reference valid entities
        entity_ids = {e.get("name") for e in self.ctx.working_model.get("entities", [])}
        for rel in self.ctx.working_model.get("relationships", []):
            if rel.get("firstEntity") not in entity_ids:
                warnings.append(f"Relationship references unknown entity: {rel.get('firstEntity')}")
            if rel.get("secondEntity") not in entity_ids:
                warnings.append(f"Relationship references unknown entity: {rel.get('secondEntity')}")

        return ToolResult(success=len(errors) == 0, data={"errors": errors, "warnings": warnings})

    def finish(self, summary: str) -> ToolResult:
        """
        Complete the generation, returning the final model.

        Args:
            summary: Summary of what was generated

        Returns:
            ToolResult with final model and coverage
        """
        self.ctx.status = "completed"

        coverage_result = self.coverage_report()

        return ToolResult(
            success=True,
            data={
                "summary": summary,
                "model": self.ctx.working_model,
                "coverage": coverage_result.data if coverage_result.success else None,
            },
        )
