"""
Governor Agent Tools

Tools for the Governor Agent responsible for data quality, SLA enforcement, and lineage tracking.

Tool Categories:
1. QualityCheckTools: Data quality checking
2. LineageTools: Data lineage tracking
3. ContractTools: Data contract/SLA management
4. MetadataTools: Metadata management

This is a standalone implementation that doesn't require external governance modules.
"""

from __future__ import annotations

import logging
import uuid
from dataclasses import dataclass, field
from datetime import datetime, timezone
from enum import Enum
from typing import Any

logger = logging.getLogger(__name__)


# ============================================================================
# Types for Governance Tools
# ============================================================================


class QualityDimension(str, Enum):
    """Data quality dimensions."""
    COMPLETENESS = "completeness"
    ACCURACY = "accuracy"
    TIMELINESS = "timeliness"
    VALIDITY = "validity"
    UNIQUENESS = "uniqueness"


class ContractStatus(str, Enum):
    """Contract status."""
    ACTIVE = "active"
    DRAFT = "draft"
    DEPRECATED = "deprecated"


class LineageNodeType(str, Enum):
    """Lineage node types."""
    TABLE = "table"
    VIEW = "view"
    PROCESS = "process"
    EXTERNAL = "external"


@dataclass
class QualityCheckResult:
    """Result of a quality check."""
    dimension: QualityDimension
    table_name: str
    score: float  # 0.0 to 1.0
    passed: bool
    issues: list[str] = field(default_factory=list)
    checked_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))


@dataclass
class LineageNode:
    """A node in the data lineage graph."""
    node_id: str
    name: str
    node_type: LineageNodeType
    metadata: dict[str, Any] = field(default_factory=dict)


@dataclass
class LineageEdge:
    """An edge in the data lineage graph."""
    source_id: str
    target_id: str
    relationship: str = "derives_from"


@dataclass
class LineageGraph:
    """Data lineage graph."""
    graph_id: str
    name: str
    nodes: list[LineageNode] = field(default_factory=list)
    edges: list[LineageEdge] = field(default_factory=list)
    created_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))


@dataclass
class ContractSLA:
    """SLA for a data contract."""
    max_latency_hours: float = 24.0
    min_completeness: float = 0.95
    min_accuracy: float = 0.99
    max_null_ratio: float = 0.05


@dataclass
class DataContract:
    """Data contract definition."""
    contract_id: str
    name: str
    table_name: str
    owner: str
    sla: ContractSLA
    status: ContractStatus = ContractStatus.ACTIVE
    created_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))


@dataclass
class TableMetadata:
    """Metadata for a table."""
    table_name: str
    schema_name: str = "public"
    description: str = ""
    owner: str = ""
    tags: list[str] = field(default_factory=list)
    columns: dict[str, str] = field(default_factory=dict)  # column_name -> type
    created_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
    updated_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))


# ============================================================================
# Quality Check Tools
# ============================================================================


class QualityCheckTools:
    """Tools for data quality checking"""

    def __init__(self):
        self._rules: dict[str, list[dict]] = {}
        self._results: list[QualityCheckResult] = []
        logger.info("QualityCheckTools initialized")

    def list_quality_rules(self) -> list[dict[str, Any]]:
        """List all registered quality rules."""
        rules = []
        for table, table_rules in self._rules.items():
            for rule in table_rules:
                rules.append({
                    "table_name": table,
                    **rule,
                })
        return rules

    def register_rule(
        self,
        table_name: str,
        dimension: str,
        rule_type: str,
        threshold: float = 0.95,
        **kwargs,
    ) -> dict[str, Any]:
        """Register a new quality rule."""
        rule = {
            "rule_id": str(uuid.uuid4())[:8],
            "dimension": dimension,
            "rule_type": rule_type,
            "threshold": threshold,
            **kwargs,
        }
        if table_name not in self._rules:
            self._rules[table_name] = []
        self._rules[table_name].append(rule)
        return rule

    def check_data_quality(
        self,
        table_name: str,
        dimension: str | None = None,
    ) -> list[dict[str, Any]]:
        """
        Run quality checks on a table.

        Returns mock results since this is standalone.
        """
        dimensions = [QualityDimension(dimension)] if dimension else list(QualityDimension)
        results = []

        for dim in dimensions:
            result = QualityCheckResult(
                dimension=dim,
                table_name=table_name,
                score=0.98,  # Mock high score
                passed=True,
                issues=[],
            )
            self._results.append(result)
            results.append({
                "dimension": dim.value,
                "table_name": table_name,
                "score": result.score,
                "passed": result.passed,
                "issues": result.issues,
                "checked_at": result.checked_at.isoformat(),
            })

        return results

    def get_quality_statistics(self) -> dict[str, Any]:
        """Get quality statistics summary."""
        if not self._results:
            return {
                "total_checks": 0,
                "passed": 0,
                "failed": 0,
                "average_score": 0.0,
            }

        passed = sum(1 for r in self._results if r.passed)
        avg_score = sum(r.score for r in self._results) / len(self._results)

        return {
            "total_checks": len(self._results),
            "passed": passed,
            "failed": len(self._results) - passed,
            "average_score": avg_score,
        }

    def get_predefined_rule_sets(self) -> dict[str, Any]:
        """Get predefined rule sets for common data types."""
        return {
            "bsr_rules": {
                "description": "Amazon BSR data quality rules",
                "rules": [
                    {"dimension": "completeness", "min_threshold": 0.95},
                    {"dimension": "timeliness", "max_hours": 24},
                ],
            },
            "content_rules": {
                "description": "Content data quality rules",
                "rules": [
                    {"dimension": "completeness", "min_threshold": 0.99},
                    {"dimension": "accuracy", "min_threshold": 0.95},
                ],
            },
        }


# ============================================================================
# Lineage Tools
# ============================================================================


class LineageTools:
    """Tools for data lineage tracking"""

    def __init__(self):
        self._graphs: dict[str, LineageGraph] = {}
        logger.info("LineageTools initialized")

    def create_lineage_graph(self, name: str) -> dict[str, Any]:
        """Create a new lineage graph."""
        graph_id = str(uuid.uuid4())[:8]
        graph = LineageGraph(
            graph_id=graph_id,
            name=name,
        )
        self._graphs[graph_id] = graph
        return {
            "graph_id": graph_id,
            "name": name,
            "created_at": graph.created_at.isoformat(),
        }

    def add_lineage_node(
        self,
        graph_id: str,
        name: str,
        node_type: str,
        metadata: dict[str, Any] | None = None,
    ) -> dict[str, Any]:
        """Add a node to a lineage graph."""
        if graph_id not in self._graphs:
            return {"error": f"Graph {graph_id} not found"}

        node_id = str(uuid.uuid4())[:8]
        node = LineageNode(
            node_id=node_id,
            name=name,
            node_type=LineageNodeType(node_type),
            metadata=metadata or {},
        )
        self._graphs[graph_id].nodes.append(node)

        return {
            "node_id": node_id,
            "name": name,
            "node_type": node_type,
        }

    def link_nodes(
        self,
        graph_id: str,
        source_id: str,
        target_id: str,
        relationship: str = "derives_from",
    ) -> dict[str, Any]:
        """Link two nodes in a lineage graph."""
        if graph_id not in self._graphs:
            return {"error": f"Graph {graph_id} not found"}

        edge = LineageEdge(
            source_id=source_id,
            target_id=target_id,
            relationship=relationship,
        )
        self._graphs[graph_id].edges.append(edge)

        return {
            "source_id": source_id,
            "target_id": target_id,
            "relationship": relationship,
        }

    def get_upstream(self, node_name: str) -> list[dict[str, Any]]:
        """Get upstream dependencies of a node."""
        upstream = []
        for graph in self._graphs.values():
            for edge in graph.edges:
                if edge.target_id == node_name:
                    upstream.append({
                        "source": edge.source_id,
                        "relationship": edge.relationship,
                    })
        return upstream

    def get_downstream(self, node_name: str) -> list[dict[str, Any]]:
        """Get downstream dependents of a node."""
        downstream = []
        for graph in self._graphs.values():
            for edge in graph.edges:
                if edge.source_id == node_name:
                    downstream.append({
                        "target": edge.target_id,
                        "relationship": edge.relationship,
                    })
        return downstream

    def impact_analysis(self, node_name: str) -> dict[str, Any]:
        """Analyze the impact of changes to a node."""
        downstream = self.get_downstream(node_name)
        return {
            "node_name": node_name,
            "direct_impacts": len(downstream),
            "downstream_nodes": downstream,
            "risk_level": "low" if len(downstream) < 5 else "medium" if len(downstream) < 10 else "high",
        }

    def list_lineage_graphs(self) -> list[dict[str, Any]]:
        """List all lineage graphs."""
        return [
            {
                "graph_id": graph.graph_id,
                "name": graph.name,
                "nodes_count": len(graph.nodes),
                "edges_count": len(graph.edges),
                "created_at": graph.created_at.isoformat(),
            }
            for graph in self._graphs.values()
        ]


# ============================================================================
# Contract Tools
# ============================================================================


class ContractTools:
    """Tools for data contract management"""

    def __init__(self):
        self._contracts: dict[str, DataContract] = {}
        self._validations: list[dict] = []
        logger.info("ContractTools initialized")

    def register_contract(
        self,
        name: str,
        table_name: str,
        owner: str,
        max_latency_hours: float = 24.0,
        min_completeness: float = 0.95,
        min_accuracy: float = 0.99,
    ) -> dict[str, Any]:
        """Register a new data contract."""
        contract_id = str(uuid.uuid4())[:8]
        sla = ContractSLA(
            max_latency_hours=max_latency_hours,
            min_completeness=min_completeness,
            min_accuracy=min_accuracy,
        )
        contract = DataContract(
            contract_id=contract_id,
            name=name,
            table_name=table_name,
            owner=owner,
            sla=sla,
        )
        self._contracts[contract_id] = contract

        return {
            "contract_id": contract_id,
            "name": name,
            "table_name": table_name,
            "owner": owner,
            "status": contract.status.value,
        }

    def validate_against_contract(
        self,
        contract_id: str,
    ) -> dict[str, Any]:
        """Validate data against a contract."""
        if contract_id not in self._contracts:
            return {"error": f"Contract {contract_id} not found"}

        contract = self._contracts[contract_id]
        # Mock validation result
        result = {
            "contract_id": contract_id,
            "contract_name": contract.name,
            "table_name": contract.table_name,
            "compliant": True,
            "checks": [
                {"check": "latency", "passed": True, "value": 12.0, "threshold": contract.sla.max_latency_hours},
                {"check": "completeness", "passed": True, "value": 0.98, "threshold": contract.sla.min_completeness},
                {"check": "accuracy", "passed": True, "value": 0.995, "threshold": contract.sla.min_accuracy},
            ],
            "validated_at": datetime.now(timezone.utc).isoformat(),
        }
        self._validations.append(result)
        return result

    def list_contracts(self) -> list[dict[str, Any]]:
        """List all contracts."""
        return [
            {
                "contract_id": c.contract_id,
                "name": c.name,
                "table_name": c.table_name,
                "owner": c.owner,
                "status": c.status.value,
                "sla": {
                    "max_latency_hours": c.sla.max_latency_hours,
                    "min_completeness": c.sla.min_completeness,
                    "min_accuracy": c.sla.min_accuracy,
                },
            }
            for c in self._contracts.values()
        ]

    def activate_contract(self, contract_id: str) -> dict[str, Any]:
        """Activate a contract."""
        if contract_id not in self._contracts:
            return {"error": f"Contract {contract_id} not found"}

        self._contracts[contract_id].status = ContractStatus.ACTIVE
        return {
            "contract_id": contract_id,
            "status": "active",
        }

    def get_contract_stats(self) -> dict[str, Any]:
        """Get contract statistics."""
        active = sum(1 for c in self._contracts.values() if c.status == ContractStatus.ACTIVE)
        return {
            "total_contracts": len(self._contracts),
            "active": active,
            "draft": len(self._contracts) - active,
            "validations_run": len(self._validations),
            "compliance_rate": 1.0 if not self._validations else (
                sum(1 for v in self._validations if v.get("compliant")) / len(self._validations)
            ),
        }


# ============================================================================
# Metadata Tools
# ============================================================================


class MetadataTools:
    """Tools for metadata management"""

    def __init__(self):
        self._metadata: dict[str, TableMetadata] = {}
        logger.info("MetadataTools initialized")

    def register_schema(
        self,
        table_name: str,
        schema_name: str = "public",
        description: str = "",
        columns: dict[str, str] | None = None,
    ) -> dict[str, Any]:
        """Register table schema metadata."""
        metadata = TableMetadata(
            table_name=table_name,
            schema_name=schema_name,
            description=description,
            columns=columns or {},
        )
        key = f"{schema_name}.{table_name}"
        self._metadata[key] = metadata

        return {
            "table_name": table_name,
            "schema_name": schema_name,
            "registered": True,
        }

    def add_tag(self, table_name: str, tag: str, schema_name: str = "public") -> dict[str, Any]:
        """Add a tag to a table."""
        key = f"{schema_name}.{table_name}"
        if key not in self._metadata:
            self._metadata[key] = TableMetadata(table_name=table_name, schema_name=schema_name)

        if tag not in self._metadata[key].tags:
            self._metadata[key].tags.append(tag)

        return {
            "table_name": table_name,
            "tag": tag,
            "added": True,
        }

    def set_owner(self, table_name: str, owner: str, schema_name: str = "public") -> dict[str, Any]:
        """Set the owner of a table."""
        key = f"{schema_name}.{table_name}"
        if key not in self._metadata:
            self._metadata[key] = TableMetadata(table_name=table_name, schema_name=schema_name)

        self._metadata[key].owner = owner
        self._metadata[key].updated_at = datetime.now(timezone.utc)

        return {
            "table_name": table_name,
            "owner": owner,
            "set": True,
        }

    def get_metadata(self, table_name: str, schema_name: str = "public") -> dict[str, Any]:
        """Get metadata for a table."""
        key = f"{schema_name}.{table_name}"
        if key not in self._metadata:
            return {"error": f"No metadata found for {key}"}

        meta = self._metadata[key]
        return {
            "table_name": meta.table_name,
            "schema_name": meta.schema_name,
            "description": meta.description,
            "owner": meta.owner,
            "tags": meta.tags,
            "columns": meta.columns,
            "created_at": meta.created_at.isoformat(),
            "updated_at": meta.updated_at.isoformat(),
        }

    def get_statistics(self) -> dict[str, Any]:
        """Get metadata statistics."""
        return {
            "total_tables": len(self._metadata),
            "tables_with_owner": sum(1 for m in self._metadata.values() if m.owner),
            "tables_with_description": sum(1 for m in self._metadata.values() if m.description),
            "total_tags": sum(len(m.tags) for m in self._metadata.values()),
        }


# ============================================================================
# Tool Definitions for Claude Agent SDK
# ============================================================================


GOVERNOR_TOOLS = [
    # Quality Check Tools
    {
        "name": "check_data_quality",
        "description": "Run data quality checks on a table. Checks completeness, accuracy, timeliness, validity, and uniqueness.",
        "input_schema": {
            "type": "object",
            "properties": {
                "table_name": {
                    "type": "string",
                    "description": "Name of the table to check",
                },
                "dimension": {
                    "type": "string",
                    "description": "Specific dimension to check (completeness, accuracy, timeliness, validity, uniqueness). If not specified, checks all.",
                    "enum": ["completeness", "accuracy", "timeliness", "validity", "uniqueness"],
                },
            },
            "required": ["table_name"],
        },
    },
    {
        "name": "list_quality_rules",
        "description": "List all registered data quality rules.",
        "input_schema": {"type": "object", "properties": {}},
    },
    {
        "name": "get_quality_statistics",
        "description": "Get summary statistics of data quality checks.",
        "input_schema": {"type": "object", "properties": {}},
    },
    {
        "name": "get_predefined_rule_sets",
        "description": "Get predefined rule sets for common data types like BSR and Content.",
        "input_schema": {"type": "object", "properties": {}},
    },
    # Lineage Tools
    {
        "name": "create_lineage_graph",
        "description": "Create a new data lineage graph.",
        "input_schema": {
            "type": "object",
            "properties": {
                "name": {"type": "string", "description": "Name of the lineage graph"},
            },
            "required": ["name"],
        },
    },
    {
        "name": "add_lineage_node",
        "description": "Add a node to a lineage graph.",
        "input_schema": {
            "type": "object",
            "properties": {
                "graph_id": {"type": "string", "description": "ID of the graph"},
                "name": {"type": "string", "description": "Name of the node"},
                "node_type": {"type": "string", "description": "Type of node", "enum": ["table", "view", "process", "external"]},
                "metadata": {"type": "object", "description": "Additional metadata"},
            },
            "required": ["graph_id", "name", "node_type"],
        },
    },
    {
        "name": "link_nodes",
        "description": "Link two nodes in a lineage graph.",
        "input_schema": {
            "type": "object",
            "properties": {
                "graph_id": {"type": "string", "description": "ID of the graph"},
                "source_id": {"type": "string", "description": "Source node ID"},
                "target_id": {"type": "string", "description": "Target node ID"},
                "relationship": {"type": "string", "description": "Relationship type", "default": "derives_from"},
            },
            "required": ["graph_id", "source_id", "target_id"],
        },
    },
    {
        "name": "get_upstream",
        "description": "Get upstream dependencies of a data asset.",
        "input_schema": {
            "type": "object",
            "properties": {
                "node_name": {"type": "string", "description": "Name of the node to trace upstream"},
            },
            "required": ["node_name"],
        },
    },
    {
        "name": "get_downstream",
        "description": "Get downstream dependents of a data asset.",
        "input_schema": {
            "type": "object",
            "properties": {
                "node_name": {"type": "string", "description": "Name of the node to trace downstream"},
            },
            "required": ["node_name"],
        },
    },
    {
        "name": "impact_analysis",
        "description": "Analyze the impact of changes to a data asset.",
        "input_schema": {
            "type": "object",
            "properties": {
                "node_name": {"type": "string", "description": "Name of the node to analyze"},
            },
            "required": ["node_name"],
        },
    },
    {
        "name": "list_lineage_graphs",
        "description": "List all data lineage graphs.",
        "input_schema": {"type": "object", "properties": {}},
    },
    # Contract Tools
    {
        "name": "register_contract",
        "description": "Register a new data contract with SLA.",
        "input_schema": {
            "type": "object",
            "properties": {
                "name": {"type": "string", "description": "Contract name"},
                "table_name": {"type": "string", "description": "Table the contract applies to"},
                "owner": {"type": "string", "description": "Contract owner"},
                "max_latency_hours": {"type": "number", "description": "Maximum acceptable data latency in hours", "default": 24.0},
                "min_completeness": {"type": "number", "description": "Minimum completeness score (0-1)", "default": 0.95},
                "min_accuracy": {"type": "number", "description": "Minimum accuracy score (0-1)", "default": 0.99},
            },
            "required": ["name", "table_name", "owner"],
        },
    },
    {
        "name": "validate_against_contract",
        "description": "Validate data against a registered contract.",
        "input_schema": {
            "type": "object",
            "properties": {
                "contract_id": {"type": "string", "description": "ID of the contract to validate against"},
            },
            "required": ["contract_id"],
        },
    },
    {
        "name": "list_contracts",
        "description": "List all data contracts.",
        "input_schema": {"type": "object", "properties": {}},
    },
    {
        "name": "activate_contract",
        "description": "Activate a data contract.",
        "input_schema": {
            "type": "object",
            "properties": {
                "contract_id": {"type": "string", "description": "ID of the contract to activate"},
            },
            "required": ["contract_id"],
        },
    },
    {
        "name": "get_contract_stats",
        "description": "Get contract statistics.",
        "input_schema": {"type": "object", "properties": {}},
    },
    # Metadata Tools
    {
        "name": "register_schema",
        "description": "Register table schema metadata.",
        "input_schema": {
            "type": "object",
            "properties": {
                "table_name": {"type": "string", "description": "Table name"},
                "schema_name": {"type": "string", "description": "Schema name", "default": "public"},
                "description": {"type": "string", "description": "Table description"},
                "columns": {"type": "object", "description": "Column name to type mapping"},
            },
            "required": ["table_name"],
        },
    },
    {
        "name": "add_tag",
        "description": "Add a tag to a table.",
        "input_schema": {
            "type": "object",
            "properties": {
                "table_name": {"type": "string", "description": "Table name"},
                "tag": {"type": "string", "description": "Tag to add"},
                "schema_name": {"type": "string", "description": "Schema name", "default": "public"},
            },
            "required": ["table_name", "tag"],
        },
    },
    {
        "name": "set_owner",
        "description": "Set the owner of a table.",
        "input_schema": {
            "type": "object",
            "properties": {
                "table_name": {"type": "string", "description": "Table name"},
                "owner": {"type": "string", "description": "Owner name"},
                "schema_name": {"type": "string", "description": "Schema name", "default": "public"},
            },
            "required": ["table_name", "owner"],
        },
    },
    {
        "name": "get_metadata",
        "description": "Get metadata for a table.",
        "input_schema": {
            "type": "object",
            "properties": {
                "table_name": {"type": "string", "description": "Table name"},
                "schema_name": {"type": "string", "description": "Schema name", "default": "public"},
            },
            "required": ["table_name"],
        },
    },
    {
        "name": "get_statistics",
        "description": "Get metadata statistics.",
        "input_schema": {"type": "object", "properties": {}},
    },
]


__all__ = [
    "QualityCheckTools",
    "LineageTools",
    "ContractTools",
    "MetadataTools",
    "GOVERNOR_TOOLS",
    "QualityDimension",
    "ContractStatus",
    "LineageNodeType",
]
