"""
Master Orchestrator Agent with Claude Agent SDK

Autonomous agent for managing the Smart Data Factory:
- Health monitoring and decision making
- ETL orchestration
- Alert and approval handling
- Scheduled task execution

Architecture:
- Claude Agent SDK for intelligent decision-making
- Token-efficient design (~100 tokens per decision cycle)
- Human-in-the-loop for critical operations
- Supports both Anthropic and OpenAI-compatible APIs (e.g., Docker Model Runner)
"""

from __future__ import annotations

import asyncio
import json
import logging
import time
from datetime import datetime, timezone, date
from typing import Any, Literal

from .tool_logger import log_tool_call
from .types import (
    AgentDecision,
    AgentExecutionResult,
    AgentPriority,
    ApprovalStatus,
    HealthStatus,
    SystemHealth,
    TriggerType,
    HUMAN_APPROVAL_REQUIRED,
)
from .tools import (
    HealthCheckTools,
    ETLExecutorTools,
    AlertSenderTools,
    ALL_TOOLS,
)

logger = logging.getLogger(__name__)


class MasterOrchestratorAgent:
    """
    Master Orchestrator Agent - Autonomous Data Factory Manager

    Uses Claude Agent SDK for intelligent decision-making:
    1. Monitors system health
    2. Makes decisions about remediation
    3. Executes actions (with human approval when needed)
    4. Reports results

    Token consumption: ~100 tokens per decision cycle
    - Health check context: ~50 tokens
    - Decision reasoning: ~50 tokens
    """

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

## Your Role
You are responsible for maintaining the health and efficiency of the data factory.
Your goal is to ensure all data pipelines run smoothly and data quality is maintained.

## Available Tools
1. **check_pipeline_health** - Check status of all ETL pipelines
2. **check_layer_health** - Check status of data warehouse layers (ODS/DWD/DWS/ADS)
3. **get_system_health** - Get comprehensive system health report with recommendations
4. **run_etl** - Run a specific ETL module (ga4, bsr, visibility, etc.)
5. **send_alert** - Send alert notifications
6. **request_human_approval** - Request human approval for critical actions

## Decision Rules
1. **Stale Data (> 24h)**: Run the corresponding ETL module
2. **Critical Data (> 48h)**: Run ETL with high priority + send alert
3. **Empty Layers**: Investigate and report, may need manual intervention
4. **Quality Issues (< 95%)**: Run governance checks + alert
5. **Destructive Actions**: Always request human approval first

## Safety Red Lines (ALWAYS require human approval)
- DELETE operations
- TRUNCATE operations
- DROP table/schema
- Rollback > 24 hours of data
- New data source integration
- SLA threshold modifications
- Actions impacting > 3 downstream systems
- Cost exceeding 150% of budget

## Response Format
1. First, analyze the current system state
2. Identify any issues that need attention
3. Propose actions with clear reasoning
4. Execute safe actions immediately
5. Request approval for critical actions
6. Report results clearly

Be concise but thorough. Prioritize system stability."""

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

        Args:
            api_key: API key (uses env if not provided for Anthropic)
            model: Model to use (e.g., "claude-sonnet-4-20250514" or "ai/gpt-oss")
            db_url: Database connection URL
            slack_webhook_url: Slack webhook for alerts
            provider: API provider ("anthropic" or "openai" for OpenAI-compatible APIs)
            base_url: Base URL for OpenAI-compatible APIs (e.g., "http://localhost:12434/v1")
        """
        self.provider = provider
        self.model = model
        self.base_url = base_url

        if provider == "anthropic":
            from anthropic import Anthropic
            self.client = Anthropic(api_key=api_key) if api_key else Anthropic()
        else:
            # OpenAI-compatible API (Docker Model Runner, etc.)
            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.health_tools = HealthCheckTools(db_url=db_url)
        self.etl_tools = ETLExecutorTools()
        self.alert_tools = AlertSenderTools(slack_webhook_url=slack_webhook_url)

        # Conversation history for agent loop
        self._messages: list[dict[str, Any]] = []

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

        # Stats
        self._stats = {
            "total_runs": 0,
            "successful_runs": 0,
            "failed_runs": 0,
            "decisions_made": 0,
            "approvals_requested": 0,
            "etls_triggered": 0,
            "alerts_sent": 0,
            "total_tokens": 0,
        }

    async def close(self):
        """Clean up resources"""
        await self.health_tools.close()

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

    async def _execute_tool(self, tool_name: str, tool_input: dict[str, Any]) -> Any:
        """
        Execute a tool and return the result

        Token consumption: 0 (pure Python execution)
        """
        logger.info(f"Executing tool: {tool_name} with input: {tool_input}")
        start_time = time.perf_counter()
        result = None
        success = True
        error_msg = None

        try:
            # Health check tools
            if tool_name == "check_pipeline_health":
                pipelines = await self.health_tools.check_pipeline_health()
                result = [
                    {
                        "module": p.module,
                        "status": p.status.value,
                        "hours_since_run": p.hours_since_run,
                        "last_data_date": p.last_data_date,
                    }
                    for p in pipelines
                ]

            elif tool_name == "check_single_pipeline":
                pipeline = await self.health_tools.check_single_pipeline(tool_input["module"])
                result = {
                    "module": pipeline.module,
                    "status": pipeline.status.value,
                    "hours_since_run": pipeline.hours_since_run,
                    "last_data_date": pipeline.last_data_date,
                    "last_error": pipeline.last_error,
                }

            elif tool_name == "check_layer_health":
                layers = await self.health_tools.check_layer_health()
                result = [
                    {
                        "layer": l.layer,
                        "tables_count": l.tables_count,
                        "total_rows": l.total_rows,
                        "status": l.status.value,
                        "issues": l.issues,
                    }
                    for l in layers
                ]

            elif tool_name == "get_system_health":
                health = await self.health_tools.get_system_health()
                result = health.to_dict()

            # ETL tools
            elif tool_name == "run_etl":
                target_date = None
                if tool_input.get("target_date"):
                    target_date = date.fromisoformat(tool_input["target_date"])

                priority = AgentPriority(tool_input.get("priority", "medium"))
                result = await self.etl_tools.run_etl(
                    module=tool_input["module"],
                    target_date=target_date,
                    priority=priority,
                )
                self._stats["etls_triggered"] += 1

            elif tool_name == "run_multiple_etls":
                result = await self.etl_tools.run_multiple_etls(
                    modules=tool_input["modules"],
                    parallel=tool_input.get("parallel", False),
                )
                self._stats["etls_triggered"] += len(tool_input["modules"])

            elif tool_name == "list_available_modules":
                result = self.etl_tools.list_available_modules()

            # Alert tools
            elif tool_name == "send_alert":
                result = await self.alert_tools.send_alert(
                    message=tool_input["message"],
                    severity=tool_input.get("severity", "info"),
                    channel=tool_input.get("channel", "log"),
                    metadata=tool_input.get("metadata"),
                )
                self._stats["alerts_sent"] += 1

            elif tool_name == "request_human_approval":
                decision = AgentDecision(
                    action=tool_input["action"],
                    reason=tool_input["reason"],
                    priority=AgentPriority(tool_input.get("priority", "medium")),
                    requires_approval=True,
                    estimated_impact=tool_input.get("estimated_impact", ""),
                )
                approval = await self.alert_tools.request_human_approval(decision)
                self._stats["approvals_requested"] += 1
                result = {
                    "request_id": approval.request_id,
                    "status": approval.status.value,
                    "expires_at": approval.expires_at.isoformat() if approval.expires_at else None,
                }

            elif tool_name == "get_pending_approvals":
                approvals = self.alert_tools.get_pending_approvals()
                result = [
                    {
                        "request_id": a.request_id,
                        "action": a.decision.action,
                        "reason": a.decision.reason,
                        "status": a.status.value,
                    }
                    for a in approvals
                ]

            elif tool_name == "get_audit_log":
                result = self.alert_tools.get_audit_log(limit=tool_input.get("limit", 100))

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

            return result

        except Exception as e:
            logger.error(f"Tool execution failed: {tool_name} - {e}")
            success = False
            error_msg = str(e)
            result = {"error": str(e)}
            return result

        finally:
            # Log tool call to monitoring database
            duration_ms = (time.perf_counter() - start_time) * 1000
            if self._execution_id:
                asyncio.create_task(log_tool_call(
                    execution_id=self._execution_id,
                    tool_name=tool_name,
                    input_params=tool_input,
                    output_result=result,
                    success=success,
                    duration_ms=duration_ms,
                    error_message=error_msg,
                ))

    # ========================================================================
    # OpenAI Format Conversion
    # ========================================================================

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

    def _convert_messages_to_openai_format(self) -> list[dict]:
        """Convert conversation messages to OpenAI format"""
        openai_messages = [
            {"role": "system", "content": self.SYSTEM_PROMPT}
        ]

        for msg in self._messages:
            if msg["role"] == "user":
                content = msg["content"]
                if isinstance(content, str):
                    openai_messages.append({"role": "user", "content": content})
                elif isinstance(content, list):
                    # Handle tool results from Anthropic format
                    for item in content:
                        if item.get("type") == "tool_result":
                            openai_messages.append({
                                "role": "tool",
                                "tool_call_id": item.get("tool_use_id", "unknown"),
                                "content": item.get("content", ""),
                            })
                        else:
                            openai_messages.append({"role": "user", "content": str(item)})
            elif msg["role"] == "assistant":
                content = msg.get("content")
                tool_calls = msg.get("tool_calls")

                if tool_calls:
                    openai_messages.append({
                        "role": "assistant",
                        "content": content if isinstance(content, str) else None,
                        "tool_calls": tool_calls,
                    })
                elif isinstance(content, str):
                    openai_messages.append({"role": "assistant", "content": content})
                elif isinstance(content, list):
                    # Convert Anthropic content blocks to text
                    text_parts = []
                    for block in content:
                        if isinstance(block, dict) and block.get("type") == "text":
                            text_parts.append(block.get("text", ""))
                    if text_parts:
                        openai_messages.append({"role": "assistant", "content": "\n".join(text_parts)})
            elif msg["role"] == "tool":
                openai_messages.append(msg)

        return openai_messages

    # ========================================================================
    # Agent Loop
    # ========================================================================

    async def run(
        self,
        task: str | None = None,
        trigger_type: TriggerType = TriggerType.MANUAL,
        max_turns: int = 10,
    ) -> AgentExecutionResult:
        """
        Run the agent loop

        Args:
            task: Optional specific task (if None, runs health check)
            trigger_type: What triggered this run
            max_turns: Maximum conversation turns

        Returns:
            AgentExecutionResult with actions taken
        """
        start_time = datetime.now(timezone.utc)
        self._stats["total_runs"] += 1

        # Default task is health check
        if not task:
            task = "Perform a health check on the data factory. Check all pipelines and data warehouse layers. If any issues are found, take appropriate action to fix them."

        # Initialize conversation
        self._messages = [
            {"role": "user", "content": task}
        ]

        decisions_made = []
        errors = []
        total_tokens = 0

        try:
            # Agent loop
            for turn in range(max_turns):
                logger.info(f"Agent turn {turn + 1}/{max_turns}")

                if self.provider == "anthropic":
                    # Call Anthropic Claude API
                    response = self.client.messages.create(
                        model=self.model,
                        max_tokens=4096,
                        system=self.SYSTEM_PROMPT,
                        tools=ALL_TOOLS,
                        messages=self._messages,
                    )

                    # Track tokens
                    total_tokens += response.usage.input_tokens + response.usage.output_tokens

                    # Process response
                    assistant_content = []
                    tool_results = []

                    for block in response.content:
                        if block.type == "text":
                            assistant_content.append({"type": "text", "text": block.text})
                            logger.info(f"Agent: {block.text[:200]}...")

                        elif block.type == "tool_use":
                            assistant_content.append({
                                "type": "tool_use",
                                "id": block.id,
                                "name": block.name,
                                "input": block.input,
                            })

                            # Execute the tool
                            result = await self._execute_tool(block.name, block.input)

                            tool_results.append({
                                "type": "tool_result",
                                "tool_use_id": block.id,
                                "content": json.dumps(result) if not isinstance(result, str) else result,
                            })

                            # Track decision
                            if block.name in ["run_etl", "send_alert", "request_human_approval"]:
                                decisions_made.append(AgentDecision(
                                    action=block.name,
                                    reason=f"Executed with input: {block.input}",
                                    priority=AgentPriority.MEDIUM,
                                    requires_approval=block.name == "request_human_approval",
                                    parameters=block.input,
                                ))
                                self._stats["decisions_made"] += 1

                    # Add assistant message
                    self._messages.append({"role": "assistant", "content": assistant_content})

                    # Add tool results if any
                    if tool_results:
                        self._messages.append({"role": "user", "content": tool_results})

                    # Check if done
                    if response.stop_reason == "end_turn":
                        logger.info("Agent completed successfully")
                        break

                else:
                    # Call OpenAI-compatible API (Docker Model Runner, etc.)
                    openai_tools = self._convert_tools_to_openai_format(ALL_TOOLS)
                    openai_messages = self._convert_messages_to_openai_format()

                    response = self.client.chat.completions.create(
                        model=self.model,
                        max_tokens=4096,
                        messages=openai_messages,
                        tools=openai_tools if openai_tools else None,
                    )

                    # Track tokens
                    if response.usage:
                        total_tokens += response.usage.prompt_tokens + response.usage.completion_tokens

                    # Process response
                    choice = response.choices[0]
                    message = choice.message

                    # Handle text response
                    if message.content:
                        logger.info(f"Agent: {message.content[:200]}...")
                        self._messages.append({"role": "assistant", "content": message.content})

                    # Handle tool calls
                    if message.tool_calls:
                        tool_results = []
                        for tool_call in message.tool_calls:
                            tool_name = tool_call.function.name
                            tool_input = json.loads(tool_call.function.arguments)

                            # Execute the tool
                            result = await self._execute_tool(tool_name, tool_input)

                            tool_results.append({
                                "tool_call_id": tool_call.id,
                                "role": "tool",
                                "content": json.dumps(result) if not isinstance(result, str) else result,
                            })

                            # Track decision
                            if tool_name in ["run_etl", "send_alert", "request_human_approval"]:
                                decisions_made.append(AgentDecision(
                                    action=tool_name,
                                    reason=f"Executed with input: {tool_input}",
                                    priority=AgentPriority.MEDIUM,
                                    requires_approval=tool_name == "request_human_approval",
                                    parameters=tool_input,
                                ))
                                self._stats["decisions_made"] += 1

                        # Add assistant message with tool calls
                        self._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
                            ]
                        })

                        # Add tool results
                        for tr in tool_results:
                            self._messages.append(tr)

                    # Check if done
                    if choice.finish_reason == "stop" and not message.tool_calls:
                        logger.info("Agent completed successfully")
                        break

            # Calculate duration
            end_time = datetime.now(timezone.utc)
            duration_ms = (end_time - start_time).total_seconds() * 1000

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

            return AgentExecutionResult(
                success=True,
                action="health_check_and_remediation",
                trigger_type=trigger_type,
                decision=decisions_made[-1] if decisions_made else None,
                result_data={
                    "turns": turn + 1,
                    "decisions_count": len(decisions_made),
                    "decisions": [
                        {"action": d.action, "parameters": d.parameters}
                        for d in decisions_made
                    ],
                },
                duration_ms=duration_ms,
                tokens_used=total_tokens,
            )

        except Exception as e:
            logger.error(f"Agent execution failed: {e}")
            self._stats["failed_runs"] += 1
            errors.append(str(e))

            return AgentExecutionResult(
                success=False,
                action="health_check_and_remediation",
                trigger_type=trigger_type,
                errors=errors,
                duration_ms=(datetime.now(timezone.utc) - start_time).total_seconds() * 1000,
            )

    # ========================================================================
    # Convenience Methods
    # ========================================================================

    async def run_health_check(self) -> AgentExecutionResult:
        """Run a health check and take action on issues"""
        return await self.run(
            task="Perform a comprehensive health check on the data factory. "
                 "Check all ETL pipelines and data warehouse layers. "
                 "If any pipelines are stale (> 24h), run them. "
                 "Send alerts for any critical issues.",
            trigger_type=TriggerType.HEALTH_CHECK,
        )

    async def run_scheduled_check(self) -> AgentExecutionResult:
        """Run scheduled health check (cron-triggered)"""
        return await self.run(
            task="This is a scheduled health check. "
                 "Check system health and take appropriate actions. "
                 "Send a summary alert with the status.",
            trigger_type=TriggerType.SCHEDULED,
        )

    async def handle_alert(self, alert_message: str) -> AgentExecutionResult:
        """Handle an incoming alert"""
        return await self.run(
            task=f"An alert has been triggered: {alert_message}\n"
                 "Investigate the issue and take appropriate action to resolve it.",
            trigger_type=TriggerType.ALERT,
        )

    async def execute_command(self, command: str) -> AgentExecutionResult:
        """Execute a manual command"""
        return await self.run(
            task=command,
            trigger_type=TriggerType.MANUAL,
        )

    def get_stats(self) -> dict[str, Any]:
        """Get agent statistics"""
        return {
            **self._stats,
            "success_rate": (
                self._stats["successful_runs"] / max(self._stats["total_runs"], 1)
            ),
            "avg_tokens_per_run": (
                self._stats["total_tokens"] / max(self._stats["total_runs"], 1)
            ),
        }


# ============================================================================
# Scheduler (Optional - for cron-based execution)
# ============================================================================

class AgentScheduler:
    """
    Scheduler for running the agent on a schedule

    Uses asyncio for scheduling (in production, use APScheduler or similar)
    """

    def __init__(self, agent: MasterOrchestratorAgent):
        self.agent = agent
        self._running = False
        self._tasks: list[asyncio.Task] = []

    async def start(self, check_interval_hours: int = 1):
        """Start the scheduler"""
        self._running = True
        logger.info(f"Starting agent scheduler with {check_interval_hours}h interval")

        while self._running:
            try:
                # Run health check
                result = await self.agent.run_scheduled_check()
                logger.info(f"Scheduled check completed: success={result.success}")

            except Exception as e:
                logger.error(f"Scheduled check failed: {e}")

            # Wait for next interval
            await asyncio.sleep(check_interval_hours * 3600)

    def stop(self):
        """Stop the scheduler"""
        self._running = False
        for task in self._tasks:
            task.cancel()
