"""
Governor Agent - Data Governance Agent

Responsible for:
1. Data quality checking and enforcement
2. Data lineage tracking
3. SLA/Contract management
4. Metadata management

Uses Claude Agent SDK with Token Zero design - all computation in Python Action Layer.
"""

from __future__ import annotations

import asyncio
import json
import logging
import time
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Any, Literal

from .tool_logger import log_tool_call
from .types import (
    AgentDecision,
    AgentExecutionResult,
    AgentPriority,
    TriggerType,
)
from .tools.governor_tools import (
    QualityCheckTools,
    LineageTools,
    ContractTools,
    MetadataTools,
    GOVERNOR_TOOLS,
)

logger = logging.getLogger(__name__)


# ============================================================================
# Governor Agent
# ============================================================================


class GovernorAgent:
    """
    Governor Agent for data governance operations.

    Capabilities:
    - Check data quality against registered rules
    - Track data lineage (source -> transform -> storage)
    - Manage data contracts and SLAs
    - Manage metadata (tags, owners, descriptions)

    Token Zero Design:
    - Skill Layer: Intent understanding (~50 tokens)
    - Action Layer: Python computation (0 tokens)
    - MCP Layer: Database operations (0 tokens)
    """

    SYSTEM_PROMPT = """You are the Governor Agent for the SNBG Smart Data Factory.

## Your Role
You are responsible for maintaining data governance through quality checks, lineage tracking, contract management, and metadata management.

## Available Tools

### Quality Check Tools
- check_data_quality: Check data quality against rules
- list_quality_rules: List all quality rules
- get_quality_statistics: Get quality statistics

### Lineage Tools
- create_lineage_graph: Create a lineage graph
- add_lineage_node: Add a node (source/transform/storage/output)
- link_lineage_nodes: Connect nodes in lineage
- get_upstream_lineage: Trace data sources
- get_downstream_lineage: Trace data consumers
- impact_analysis: Analyze change impact
- list_lineage_graphs: List all graphs

### Contract Tools
- register_contract: Register a data contract with SLA
- validate_against_contract: Validate data against SLA
- list_contracts: List all contracts
- activate_contract: Activate a contract
- get_contract_statistics: Get contract stats

### Metadata Tools
- register_schema: Register a schema
- add_metadata_tag: Add a tag to data asset
- set_data_owner: Set data owner
- get_metadata: Get metadata for an asset
- get_metadata_statistics: Get metadata stats

## Decision Guidelines

### Quality Governance
1. Check quality for critical data paths
2. Monitor quality trends over time
3. Alert on quality degradation
4. Recommend quality improvements

### Lineage Governance
1. Track all ETL pipeline flows
2. Document data transformations
3. Enable impact analysis before changes
4. Maintain up-to-date lineage graphs

### Contract Governance
1. Ensure all critical data has contracts
2. Validate data against SLAs regularly
3. Track contract violations
4. Recommend SLA adjustments

### Metadata Governance
1. Ensure all data assets have owners
2. Apply consistent tagging
3. Maintain schema documentation
4. Track metadata completeness

## Response Format
When completing a governance task, provide:
1. Summary of actions taken
2. Current governance state
3. Any issues or recommendations

Be concise and actionable in your responses.
"""

    def __init__(
        self,
        api_key: str | None = None,
        model: str = "claude-sonnet-4-20250514",
        provider: Literal["anthropic", "openai"] = "anthropic",
        base_url: str | None = None,
    ):
        """
        Initialize Governor Agent.

        Args:
            api_key: API key (Anthropic or OpenAI)
            model: Model to use
            provider: API provider ('anthropic' or 'openai')
            base_url: Base URL for OpenAI-compatible APIs
        """
        self.model = model
        self.provider = provider
        self.base_url = base_url

        # Initialize client based on provider
        if provider == "anthropic":
            from anthropic import Anthropic

            self.client = Anthropic(api_key=api_key) if api_key else Anthropic()
        else:
            from openai import OpenAI

            self.client = OpenAI(
                api_key=api_key or "local-dev-key",
                base_url=base_url or "http://localhost:12434/v1",
            )

        # Initialize tools
        self._quality_tools = QualityCheckTools()
        self._lineage_tools = LineageTools()
        self._contract_tools = ContractTools()
        self._metadata_tools = MetadataTools()

        # Current execution ID for tool logging
        self._execution_id: str | None = None

        # Execution stats
        self._stats = {
            "total_runs": 0,
            "successful_runs": 0,
            "quality_checks": 0,
            "lineage_operations": 0,
            "contract_operations": 0,
            "metadata_operations": 0,
        }

        logger.info("GovernorAgent initialized")

    # =========================================================================
    # Tool Execution
    # =========================================================================

    def _execute_tool(self, name: str, input_data: dict) -> str:
        """Execute a tool and return JSON result."""
        start_time = time.perf_counter()
        result = None
        success = True
        error_msg = None

        try:
            # Quality tools
            if name == "check_data_quality":
                result = self._quality_tools.check_data_quality(**input_data)
                self._stats["quality_checks"] += 1
            elif name == "list_quality_rules":
                result = self._quality_tools.list_quality_rules(**input_data)
            elif name == "get_quality_statistics":
                result = self._quality_tools.get_quality_statistics()

            # Lineage tools
            elif name == "create_lineage_graph":
                result = self._lineage_tools.create_lineage_graph(**input_data)
                self._stats["lineage_operations"] += 1
            elif name == "add_lineage_node":
                result = self._lineage_tools.add_lineage_node(**input_data)
                self._stats["lineage_operations"] += 1
            elif name == "link_lineage_nodes":
                result = self._lineage_tools.link_nodes(**input_data)
                self._stats["lineage_operations"] += 1
            elif name == "get_upstream_lineage":
                result = self._lineage_tools.get_upstream(**input_data)
            elif name == "get_downstream_lineage":
                result = self._lineage_tools.get_downstream(**input_data)
            elif name == "impact_analysis":
                result = self._lineage_tools.impact_analysis(**input_data)
            elif name == "list_lineage_graphs":
                result = self._lineage_tools.list_lineage_graphs()

            # Contract tools
            elif name == "register_contract":
                result = self._contract_tools.register_contract(**input_data)
                self._stats["contract_operations"] += 1
            elif name == "validate_against_contract":
                result = self._contract_tools.validate_against_contract(**input_data)
                self._stats["contract_operations"] += 1
            elif name == "list_contracts":
                result = self._contract_tools.list_contracts(**input_data)
            elif name == "activate_contract":
                result = self._contract_tools.activate_contract(**input_data)
                self._stats["contract_operations"] += 1
            elif name == "get_contract_statistics":
                result = self._contract_tools.get_contract_stats()

            # Metadata tools
            elif name == "register_schema":
                result = self._metadata_tools.register_schema(**input_data)
                self._stats["metadata_operations"] += 1
            elif name == "add_metadata_tag":
                result = self._metadata_tools.add_tag(**input_data)
                self._stats["metadata_operations"] += 1
            elif name == "set_data_owner":
                result = self._metadata_tools.set_owner(**input_data)
                self._stats["metadata_operations"] += 1
            elif name == "get_metadata":
                result = self._metadata_tools.get_metadata(**input_data)
            elif name == "get_metadata_statistics":
                result = self._metadata_tools.get_statistics()

            else:
                result = {"success": False, "error": f"Unknown tool: {name}"}
                success = False

            return json.dumps(result, default=str)

        except Exception as e:
            logger.error(f"Tool execution error: {name} - {e}")
            success = False
            error_msg = str(e)
            result = {"success": False, "error": str(e)}
            return json.dumps(result, default=str)

        finally:
            # Log tool call to monitoring database
            duration_ms = (time.perf_counter() - start_time) * 1000
            if self._execution_id:
                try:
                    loop = asyncio.get_event_loop()
                    if loop.is_running():
                        asyncio.ensure_future(log_tool_call(
                            execution_id=self._execution_id,
                            tool_name=name,
                            input_params=input_data,
                            output_result=result,
                            success=success,
                            duration_ms=duration_ms,
                            error_message=error_msg,
                        ))
                except RuntimeError:
                    # No event loop running, skip logging
                    pass

    # =========================================================================
    # Anthropic API Methods
    # =========================================================================

    def _run_anthropic_loop(
        self,
        task: str,
        max_turns: int = 10,
    ) -> AgentExecutionResult:
        """Run agent loop using Anthropic API."""
        start_time = time.time()
        decisions: list[AgentDecision] = []
        errors: list[str] = []
        total_tokens = 0

        messages = [{"role": "user", "content": task}]

        for turn in range(max_turns):
            logger.info(f"Governor turn {turn + 1}/{max_turns}")

            try:
                response = self.client.messages.create(
                    model=self.model,
                    max_tokens=4096,
                    system=self.SYSTEM_PROMPT,
                    tools=GOVERNOR_TOOLS,
                    messages=messages,
                )

                total_tokens += response.usage.input_tokens + response.usage.output_tokens

                # Process response
                assistant_content = []
                has_tool_use = False

                for block in response.content:
                    if block.type == "text":
                        assistant_content.append({"type": "text", "text": block.text})
                    elif block.type == "tool_use":
                        has_tool_use = True
                        assistant_content.append({
                            "type": "tool_use",
                            "id": block.id,
                            "name": block.name,
                            "input": block.input,
                        })

                        # Execute tool
                        logger.info(f"Executing tool: {block.name} with input: {block.input}")
                        tool_result = self._execute_tool(block.name, block.input)

                        decisions.append(AgentDecision(
                            action=block.name,
                            reason="Tool execution for governance",
                            priority=AgentPriority.MEDIUM,
                            requires_approval=False,
                            parameters=block.input,
                        ))

                        # Add tool result to messages
                        messages.append({"role": "assistant", "content": assistant_content})
                        messages.append({
                            "role": "user",
                            "content": [{
                                "type": "tool_result",
                                "tool_use_id": block.id,
                                "content": tool_result,
                            }],
                        })

                # Check for completion
                if response.stop_reason == "end_turn" or not has_tool_use:
                    final_text = ""
                    for block in response.content:
                        if block.type == "text":
                            final_text = block.text
                            break

                    logger.info(f"Governor: {final_text[:200]}...")
                    logger.info("Governor completed successfully")

                    self._stats["total_runs"] += 1
                    self._stats["successful_runs"] += 1

                    return AgentExecutionResult(
                        success=True,
                        action="check_governance_status",
                        trigger_type=TriggerType.HEALTH_CHECK,
                        duration_ms=(time.time() - start_time) * 1000,
                        tokens_used=total_tokens,
                        decision=decisions[-1] if decisions else None,
                        errors=[],
                        result_data={
                            "turns": turn + 1,
                            "decisions_count": len(decisions),
                            "decisions": [
                                {"action": d.action, "parameters": d.parameters}
                                for d in decisions
                            ],
                            "final_response": final_text,
                        },
                    )

                if not has_tool_use:
                    messages.append({"role": "assistant", "content": assistant_content})

            except Exception as e:
                logger.error(f"Governor turn error: {e}")
                errors.append(str(e))
                break

        # Max turns reached
        self._stats["total_runs"] += 1

        return AgentExecutionResult(
            success=False,
            action="check_governance_status",
            trigger_type=TriggerType.HEALTH_CHECK,
            duration_ms=(time.time() - start_time) * 1000,
            tokens_used=total_tokens,
            decision=decisions[-1] if decisions else None,
            errors=errors or ["Max turns reached"],
            result_data={"turns": max_turns, "decisions_count": len(decisions)},
        )

    # =========================================================================
    # OpenAI-Compatible API Methods
    # =========================================================================

    def _convert_tools_to_openai_format(self) -> list[dict]:
        """Convert tools to OpenAI function calling format."""
        openai_tools = []
        for tool in GOVERNOR_TOOLS:
            openai_tools.append({
                "type": "function",
                "function": {
                    "name": tool["name"],
                    "description": tool["description"],
                    "parameters": tool["input_schema"],
                },
            })
        return openai_tools

    def _run_openai_loop(
        self,
        task: str,
        max_turns: int = 10,
    ) -> AgentExecutionResult:
        """Run agent loop using OpenAI-compatible API."""
        start_time = time.time()
        decisions: list[AgentDecision] = []
        errors: list[str] = []
        total_tokens = 0

        messages = [
            {"role": "system", "content": self.SYSTEM_PROMPT},
            {"role": "user", "content": task},
        ]

        openai_tools = self._convert_tools_to_openai_format()

        for turn in range(max_turns):
            logger.info(f"Governor turn {turn + 1}/{max_turns}")

            try:
                response = self.client.chat.completions.create(
                    model=self.model,
                    max_tokens=4096,
                    messages=messages,
                    tools=openai_tools,
                    tool_choice="auto",
                )

                if response.usage:
                    total_tokens += response.usage.total_tokens

                message = response.choices[0].message
                finish_reason = response.choices[0].finish_reason

                # Check for tool calls
                if message.tool_calls:
                    # Add assistant message with tool calls
                    messages.append({
                        "role": "assistant",
                        "content": message.content,
                        "tool_calls": [
                            {
                                "id": tc.id,
                                "type": "function",
                                "function": {
                                    "name": tc.function.name,
                                    "arguments": tc.function.arguments,
                                },
                            }
                            for tc in message.tool_calls
                        ],
                    })

                    # Execute each tool
                    for tool_call in message.tool_calls:
                        tool_name = tool_call.function.name
                        tool_input = json.loads(tool_call.function.arguments)

                        logger.info(f"Executing tool: {tool_name} with input: {tool_input}")
                        tool_result = self._execute_tool(tool_name, tool_input)

                        decisions.append(AgentDecision(
                            action=tool_name,
                            reason="Tool execution for governance",
                            priority=AgentPriority.MEDIUM,
                            requires_approval=False,
                            parameters=tool_input,
                        ))

                        # Add tool result
                        messages.append({
                            "role": "tool",
                            "tool_call_id": tool_call.id,
                            "content": tool_result,
                        })

                elif finish_reason == "stop":
                    # Completed
                    final_text = message.content or ""
                    logger.info(f"Governor: {final_text[:200]}...")
                    logger.info("Governor completed successfully")

                    self._stats["total_runs"] += 1
                    self._stats["successful_runs"] += 1

                    return AgentExecutionResult(
                        success=True,
                        action="check_governance_status",
                        trigger_type=TriggerType.HEALTH_CHECK,
                        duration_ms=(time.time() - start_time) * 1000,
                        tokens_used=total_tokens,
                        decision=decisions[-1] if decisions else None,
                        errors=[],
                        result_data={
                            "turns": turn + 1,
                            "decisions_count": len(decisions),
                            "decisions": [
                                {"action": d.action, "parameters": d.parameters}
                                for d in decisions
                            ],
                            "final_response": final_text,
                        },
                    )

                else:
                    # No tool calls, add message and continue
                    messages.append({
                        "role": "assistant",
                        "content": message.content,
                    })

            except Exception as e:
                logger.error(f"Governor turn error: {e}")
                errors.append(str(e))
                break

        # Max turns reached
        self._stats["total_runs"] += 1

        return AgentExecutionResult(
            success=False,
            action="check_governance_status",
            trigger_type=TriggerType.HEALTH_CHECK,
            duration_ms=(time.time() - start_time) * 1000,
            tokens_used=total_tokens,
            decision=decisions[-1] if decisions else None,
            errors=errors or ["Max turns reached"],
            result_data={"turns": max_turns, "decisions_count": len(decisions)},
        )

    # =========================================================================
    # Public API
    # =========================================================================

    async def run(
        self,
        task: str,
        max_turns: int = 10,
    ) -> AgentExecutionResult:
        """
        Run the Governor Agent with a task.

        Args:
            task: Governance task to execute
            max_turns: Maximum agent turns

        Returns:
            AgentExecutionResult
        """
        if self.provider == "anthropic":
            return self._run_anthropic_loop(task, max_turns)
        else:
            return self._run_openai_loop(task, max_turns)

    async def check_governance_status(self) -> AgentExecutionResult:
        """
        Check overall governance status.

        Returns:
            AgentExecutionResult with governance status
        """
        task = """Check the overall governance status:
1. List all quality rules and get quality statistics
2. List all lineage graphs
3. List all data contracts and get contract statistics
4. Get metadata statistics

Provide a summary of the current governance state."""

        return await self.run(task)

    async def audit_data_quality(
        self,
        data_layer: str | None = None,
    ) -> AgentExecutionResult:
        """
        Audit data quality across the system.

        Args:
            data_layer: Optional data layer to focus on (ods, dwd, dws, ads)

        Returns:
            AgentExecutionResult with audit results
        """
        layer_context = f" for the {data_layer} layer" if data_layer else ""
        task = f"""Perform a data quality audit{layer_context}:
1. List all quality rules
2. Get quality statistics
3. Identify any quality issues or trends
4. Provide recommendations for improvement"""

        return await self.run(task)

    async def analyze_lineage_impact(
        self,
        source_description: str,
    ) -> AgentExecutionResult:
        """
        Analyze impact of changes to a data source.

        Args:
            source_description: Description of the source to analyze

        Returns:
            AgentExecutionResult with impact analysis
        """
        task = f"""Analyze the impact of changes to: {source_description}

1. Find the relevant lineage graph
2. Identify the source node
3. Run impact analysis to find all downstream dependencies
4. Provide a summary of affected consumers and recommendations"""

        return await self.run(task)

    async def validate_sla_compliance(
        self,
        contract_name: str | None = None,
    ) -> AgentExecutionResult:
        """
        Validate SLA compliance for contracts.

        Args:
            contract_name: Specific contract to validate (None = all)

        Returns:
            AgentExecutionResult with compliance status
        """
        if contract_name:
            task = f"""Validate SLA compliance for contract: {contract_name}

1. Get the contract details
2. Check if data meets SLA requirements
3. Report any violations
4. Provide recommendations"""
        else:
            task = """Validate SLA compliance for all active contracts:

1. List all active contracts
2. Get contract statistics
3. Identify any contracts with violations
4. Provide compliance summary and recommendations"""

        return await self.run(task)

    def get_stats(self) -> dict:
        """Get agent statistics."""
        return {
            **self._stats,
            "success_rate": (
                self._stats["successful_runs"] / self._stats["total_runs"]
                if self._stats["total_runs"] > 0
                else 0.0
            ),
        }

    async def close(self) -> None:
        """Clean up resources."""
        logger.info("GovernorAgent closed")


# ============================================================================
# Exports
# ============================================================================

__all__ = ["GovernorAgent"]
