"""
Healer Agent - Self-Healing Autonomous Agent

Autonomous agent for system self-healing:
- Circuit breaker management
- Anomaly detection and response
- Healing action execution
- Chaos engineering tests

Architecture:
- Supports both Anthropic Claude and OpenAI-compatible APIs
- Token-efficient design
- Human-in-the-loop for critical actions
"""

from __future__ import annotations

import asyncio
import json
import logging
import time
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.healer_tools import (
    CircuitBreakerTools,
    AnomalyDetectorTools,
    HealingActionTools,
    ChaosTools,
    HEALER_TOOLS,
)

logger = logging.getLogger(__name__)


class HealerAgent:
    """
    Healer Agent - Autonomous Self-Healing System Manager

    Capabilities:
    1. Monitor circuit breaker states
    2. Detect anomalies in system metrics
    3. Execute healing actions (with approval for critical ones)
    4. Run chaos engineering experiments
    5. Report system health

    Token consumption: ~100 tokens per decision cycle
    """

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

## Your Role
You are responsible for maintaining system health through self-healing capabilities.
Your goal is to detect failures, recover automatically when safe, and escalate when needed.

## Available Tools
### Circuit Breaker Management
- **list_circuits** - List all circuit breakers and their states
- **list_open_circuits** - List all open (tripped) circuit breakers
- **get_circuit_state** - Get detailed state of a specific circuit
- **trip_circuit** - Manually trip a circuit (block requests)
- **reset_circuit** - Reset a circuit to CLOSED state

### Anomaly Detection
- **detect_anomalies** - Run anomaly detection on all metrics
- **add_metric** - Add a metric to monitor

### Healing Actions
- **execute_healing** - Execute a healing action:
  - cache_clear: Clear cache for a target
  - rate_limit: Apply rate limiting
  - circuit_open/close: Manage circuit breakers
  - retry_backlog: Retry failed items
  - alert_oncall: Alert on-call engineer
  - restart_service: Restart service (requires approval)
  - failover: Trigger failover (requires approval)
- **get_healing_history** - Get history of healing actions

### Chaos Engineering
- **run_chaos_experiment** - Run chaos test (latency/error/resource stress)
- **stop_chaos_experiment** - Stop active experiment

## Decision Rules
1. **Open Circuit (> 50% errors)**: Alert and analyze root cause
2. **High Latency (> 5s)**: Clear cache, apply rate limit
3. **Anomaly Detected**: Investigate, escalate if critical
4. **Service Failure**: Alert on-call, request approval for restart
5. **Cascading Failure**: Trip circuits to isolate, alert immediately

## Safety Red Lines (ALWAYS require human approval)
- Service restart
- Failover to backup
- Chaos experiments on production
- Circuit reset after critical failure
- Actions affecting > 3 services

## Response Format
1. Assess current system health
2. Identify issues (open circuits, anomalies)
3. Propose healing actions with reasoning
4. Execute safe actions automatically
5. Request approval for critical actions
6. Report results and recommendations

Be proactive but cautious. Prioritize system stability over aggressive recovery."""

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

        Args:
            api_key: API key (uses env if not provided for Anthropic)
            model: Model to use
            redis_url: Redis connection URL for state persistence
            provider: API provider ("anthropic" or "openai")
            base_url: Base URL for OpenAI-compatible APIs
        """
        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:
            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.circuit_tools = CircuitBreakerTools(redis_url=redis_url)
        self.anomaly_tools = AnomalyDetectorTools(redis_url=redis_url)
        self.healing_tools = HealingActionTools(redis_url=redis_url)
        self.chaos_tools = ChaosTools()

        # Conversation history
        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,
            "healings_executed": 0,
            "circuits_managed": 0,
            "anomalies_detected": 0,
            "chaos_experiments": 0,
            "total_tokens": 0,
        }

    async def initialize(self) -> None:
        """Initialize all tools"""
        await self.circuit_tools.initialize()
        await self.anomaly_tools.initialize()
        await self.healing_tools.initialize()
        logger.info("HealerAgent initialized")

    async def close(self) -> None:
        """Clean up resources"""
        await self.circuit_tools.close()
        await self.anomaly_tools.close()
        await self.healing_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"""
        logger.info(f"Executing tool: {tool_name} with input: {tool_input}")
        start_time = time.perf_counter()
        result = None
        success = True
        error_msg = None

        try:
            # Circuit Breaker Tools
            if tool_name == "list_circuits":
                result = await self.circuit_tools.list_circuits()
            elif tool_name == "list_open_circuits":
                result = await self.circuit_tools.list_open_circuits()
            elif tool_name == "get_circuit_state":
                result = await self.circuit_tools.get_circuit_state(tool_input["name"])
            elif tool_name == "trip_circuit":
                self._stats["circuits_managed"] += 1
                result = await self.circuit_tools.trip_circuit(
                    tool_input["name"],
                    tool_input.get("reason", "manual"),
                )
            elif tool_name == "reset_circuit":
                self._stats["circuits_managed"] += 1
                result = await self.circuit_tools.reset_circuit(tool_input["name"])

            # Anomaly Detection Tools
            elif tool_name == "detect_anomalies":
                result = await self.anomaly_tools.detect_anomalies()
                self._stats["anomalies_detected"] += sum(
                    len(r.get("anomalies", []))
                    for r in result.get("results", {}).values()
                )
            elif tool_name == "add_metric":
                result = await self.anomaly_tools.add_metric(
                    tool_input["metric_name"],
                    tool_input.get("detection_method", "zscore"),
                    tool_input.get("sensitivity", 3.0),
                )

            # Healing Action Tools
            elif tool_name == "execute_healing":
                self._stats["healings_executed"] += 1
                result = await self.healing_tools.execute_healing(
                    tool_input["action_type"],
                    tool_input["target"],
                    tool_input.get("parameters"),
                )
            elif tool_name == "get_healing_history":
                result = self.healing_tools.get_action_history(
                    tool_input.get("limit", 100)
                )

            # Chaos Engineering Tools
            elif tool_name == "run_chaos_experiment":
                self._stats["chaos_experiments"] += 1
                result = await self.chaos_tools.run_chaos_experiment(
                    tool_input["experiment_type"],
                    tool_input["target"],
                    tool_input.get("duration_seconds", 60),
                    tool_input.get("parameters"),
                )
            elif tool_name == "stop_chaos_experiment":
                result = self.chaos_tools.stop_experiment()

            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, tools: list[dict]) -> list[dict]:
        """Convert tool definitions to OpenAI function calling format"""
        return [
            {
                "type": "function",
                "function": {
                    "name": tool["name"],
                    "description": tool.get("description", ""),
                    "parameters": tool.get("input_schema", {"type": "object", "properties": {}}),
                }
            }
            for tool in 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):
                    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", ""),
                            })
            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):
                    text_parts = [b.get("text", "") for b in content if isinstance(b, dict) and b.get("type") == "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 healer agent loop

        Args:
            task: Optional specific task (if None, runs health assessment)
            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

        # Ensure tools are initialized
        await self.initialize()

        # Default task
        if not task:
            task = """Perform a system health assessment:
1. Check all circuit breakers for any open/tripped states
2. Run anomaly detection on monitored metrics
3. Review recent healing actions
4. If issues found, propose and execute appropriate healing actions
5. Report overall system health and any recommendations"""

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

        decisions_made = []
        errors = []
        total_tokens = 0

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

                if self.provider == "anthropic":
                    response = self.client.messages.create(
                        model=self.model,
                        max_tokens=4096,
                        system=self.SYSTEM_PROMPT,
                        tools=HEALER_TOOLS,
                        messages=self._messages,
                    )

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

                    assistant_content = []
                    tool_results = []

                    for block in response.content:
                        if block.type == "text":
                            assistant_content.append({"type": "text", "text": block.text})
                            logger.info(f"Healer: {block.text[:200]}...")
                        elif block.type == "tool_use":
                            assistant_content.append({
                                "type": "tool_use",
                                "id": block.id,
                                "name": block.name,
                                "input": block.input,
                            })

                            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,
                            })

                            if block.name in ["execute_healing", "trip_circuit", "reset_circuit"]:
                                decisions_made.append(AgentDecision(
                                    action=block.name,
                                    reason=f"Executed with input: {block.input}",
                                    priority=AgentPriority.MEDIUM,
                                    parameters=block.input,
                                ))

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

                    if tool_results:
                        self._messages.append({"role": "user", "content": tool_results})

                    if response.stop_reason == "end_turn":
                        logger.info("Healer completed successfully")
                        break

                else:
                    # OpenAI-compatible API
                    openai_tools = self._convert_tools_to_openai_format(HEALER_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,
                    )

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

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

                    if message.content:
                        logger.info(f"Healer: {message.content[:200]}...")
                        self._messages.append({"role": "assistant", "content": message.content})

                    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)

                            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,
                            })

                            if tool_name in ["execute_healing", "trip_circuit", "reset_circuit"]:
                                decisions_made.append(AgentDecision(
                                    action=tool_name,
                                    reason=f"Executed with input: {tool_input}",
                                    priority=AgentPriority.MEDIUM,
                                    parameters=tool_input,
                                ))

                        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
                            ]
                        })

                        for tr in tool_results:
                            self._messages.append(tr)

                    if choice.finish_reason == "stop" and not message.tool_calls:
                        logger.info("Healer completed successfully")
                        break

            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="self_healing_assessment",
                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"Healer execution failed: {e}")
            self._stats["failed_runs"] += 1
            errors.append(str(e))

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

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

    async def assess_health(self) -> AgentExecutionResult:
        """Run a health assessment"""
        return await self.run(
            task="Assess system health: check circuits, detect anomalies, review recent healings.",
            trigger_type=TriggerType.HEALTH_CHECK,
        )

    async def handle_failure(self, failure_description: str) -> AgentExecutionResult:
        """Handle a reported failure"""
        return await self.run(
            task=f"A failure has been reported: {failure_description}\n"
                 "Investigate and execute appropriate healing actions.",
            trigger_type=TriggerType.ALERT,
        )

    async def run_chaos_test(self, target: str, experiment_type: str = "latency_injection") -> AgentExecutionResult:
        """Run a chaos engineering test"""
        return await self.run(
            task=f"Run a chaos experiment on {target}:\n"
                 f"1. Execute {experiment_type} experiment\n"
                 "2. Monitor system response\n"
                 "3. Execute healing if needed\n"
                 "4. Report results and resilience assessment",
            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)
            ),
        }
