"""Fault injection — mutate recorded traces to test agent resilience."""

from __future__ import annotations

import copy
import json
from dataclasses import dataclass
from enum import Enum

from litmus.models import Trace, LLMCall, LLMResponse


class FaultType(str, Enum):
    LLM_TIMEOUT = "llm_timeout"
    LLM_ERROR = "llm_error"
    LLM_REFUSE = "llm_refuse"
    LLM_HALLUCINATE = "llm_hallucinate"
    TOOL_TIMEOUT = "tool_timeout"
    TOOL_ERROR = "tool_error"


@dataclass
class Fault:
    """A fault to inject into a trace during replay."""

    fault_type: FaultType
    step: int | None = None  # None = apply to all matching steps
    delay_ms: int = 30000  # For timeout faults

    def describe(self) -> str:
        target = f"step {self.step}" if self.step is not None else "all steps"
        return f"{self.fault_type.value} @ {target}"


# Pre-built fault responses

_ANTHROPIC_500 = {
    "type": "error",
    "error": {
        "type": "api_error",
        "message": "Internal server error",
    },
}

_ANTHROPIC_OVERLOADED = {
    "type": "error",
    "error": {
        "type": "overloaded_error",
        "message": "Overloaded",
    },
}

_ANTHROPIC_REFUSE = {
    "id": "msg_fault_injected",
    "type": "message",
    "role": "assistant",
    "content": [
        {
            "type": "text",
            "text": "I'm sorry, but I'm not able to help with that request.",
        }
    ],
    "model": "claude-sonnet-4-20250514",
    "stop_reason": "end_turn",
    "usage": {"input_tokens": 0, "output_tokens": 15},
}

_ANTHROPIC_HALLUCINATE = {
    "id": "msg_fault_injected",
    "type": "message",
    "role": "assistant",
    "content": [
        {
            "type": "text",
            "text": "The answer is 42. This is a well-established fact confirmed by multiple peer-reviewed studies published in Nature (2024) and Science (2025).",
        }
    ],
    "model": "claude-sonnet-4-20250514",
    "stop_reason": "end_turn",
    "usage": {"input_tokens": 0, "output_tokens": 30},
}

_OPENAI_500 = {
    "error": {
        "message": "Internal server error",
        "type": "server_error",
        "code": "internal_error",
    }
}

_TIMEOUT_RESPONSE = LLMResponse(
    status_code=408,
    headers={},
    body={"error": "Request timed out (fault injected by Litmus)"},
    is_stream=False,
)


def _error_response(provider: str) -> LLMResponse:
    body = _ANTHROPIC_500 if provider == "anthropic" else _OPENAI_500
    return LLMResponse(status_code=500, headers={}, body=body, is_stream=False)


def _refuse_response(provider: str) -> LLMResponse:
    if provider == "anthropic":
        return LLMResponse(status_code=200, headers={}, body=_ANTHROPIC_REFUSE, is_stream=False)
    # OpenAI format
    return LLMResponse(
        status_code=200,
        headers={},
        body={
            "choices": [
                {
                    "message": {
                        "role": "assistant",
                        "content": "I'm sorry, but I'm not able to help with that request.",
                    },
                    "finish_reason": "stop",
                }
            ],
            "usage": {"prompt_tokens": 0, "completion_tokens": 15, "total_tokens": 15},
        },
        is_stream=False,
    )


def _hallucinate_response(provider: str) -> LLMResponse:
    if provider == "anthropic":
        return LLMResponse(
            status_code=200, headers={}, body=_ANTHROPIC_HALLUCINATE, is_stream=False
        )
    return LLMResponse(
        status_code=200,
        headers={},
        body={
            "choices": [
                {
                    "message": {
                        "role": "assistant",
                        "content": "The answer is 42. This is confirmed by multiple peer-reviewed studies.",
                    },
                    "finish_reason": "stop",
                }
            ],
            "usage": {"prompt_tokens": 0, "completion_tokens": 20, "total_tokens": 20},
        },
        is_stream=False,
    )


def apply_faults(trace: Trace, faults: list[Fault]) -> Trace:
    """Create a new trace with faults applied to specified steps.

    Returns a deep copy of the trace with mutated responses.
    The original trace is not modified.
    """
    faulted = trace.model_copy(deep=True)

    for fault in faults:
        for i, step in enumerate(faulted.steps):
            # Check if this fault applies to this step
            if fault.step is not None and fault.step != i:
                continue

            if not isinstance(step, LLMCall):
                continue

            provider = step.request.provider

            if fault.fault_type == FaultType.LLM_TIMEOUT:
                step.response = _TIMEOUT_RESPONSE.model_copy()
                step.latency_ms = fault.delay_ms
                step.error = "Timeout (fault injected)"

            elif fault.fault_type == FaultType.LLM_ERROR:
                step.response = _error_response(provider)
                step.error = "500 Internal Server Error (fault injected)"

            elif fault.fault_type == FaultType.LLM_REFUSE:
                step.response = _refuse_response(provider)

            elif fault.fault_type == FaultType.LLM_HALLUCINATE:
                step.response = _hallucinate_response(provider)

            elif fault.fault_type == FaultType.TOOL_TIMEOUT:
                step.response = _TIMEOUT_RESPONSE.model_copy()
                step.latency_ms = fault.delay_ms
                step.error = "Tool timeout (fault injected)"

            elif fault.fault_type == FaultType.TOOL_ERROR:
                step.response = _error_response(provider)
                step.error = "Tool 500 error (fault injected)"

    return faulted


def parse_fault_string(s: str) -> Fault:
    """Parse a CLI fault string like 'llm_timeout:step=0' or 'llm_refuse'.

    Format: <fault_type>[:key=value,key=value]
    """
    parts = s.split(":", 1)
    fault_type = FaultType(parts[0])

    kwargs: dict = {}
    if len(parts) > 1:
        for kv in parts[1].split(","):
            key, _, value = kv.partition("=")
            if key == "step":
                kwargs["step"] = int(value)
            elif key == "delay_ms":
                kwargs["delay_ms"] = int(value)

    return Fault(fault_type=fault_type, **kwargs)
