"""Core failure-attribution logic for multi-agent traces.

Three methods:
  - all_at_once: feed entire conversation to the LLM
  - step_by_step: iterate step-by-step, stop at first error
  - binary_search: recursively narrow down the error location
"""

from __future__ import annotations

import random
from dataclasses import dataclass, asdict
from typing import Any

from .llm import chat, DEFAULT_MODEL, DEFAULT_MAX_TOKENS


@dataclass
class AttributionResult:
    """Result of failure attribution."""
    agent_name: str
    step_number: int
    reason: str

    def to_dict(self) -> dict:
        return asdict(self)


def _detect_agent_key(history: list[dict]) -> str:
    """Auto-detect whether agents are identified by 'name' or 'role'."""
    if history and "name" in history[0]:
        return "name"
    return "role"


def _format_history(history: list[dict], agent_key: str) -> str:
    """Format conversation history into a readable string."""
    lines = []
    for idx, entry in enumerate(history):
        agent = entry.get(agent_key, "Unknown Agent")
        content = entry.get("content", "")
        lines.append(f"Step {idx} - {agent}: {content}")
    return "\n".join(lines)


def _parse_all_at_once_result(response: str, history: list[dict], agent_key: str) -> AttributionResult:
    """Parse the LLM response from all-at-once method."""
    agent_name = "Unknown"
    step_number = -1
    reason = response

    for line in response.split("\n"):
        line_lower = line.strip().lower()
        if line_lower.startswith("agent name:"):
            agent_name = line.split(":", 1)[1].strip()
        elif line_lower.startswith("step number:"):
            try:
                step_number = int("".join(filter(str.isdigit, line.split(":", 1)[1])))
            except (ValueError, IndexError):
                pass
        elif line_lower.startswith("reason for mistake:") or line_lower.startswith("reason:"):
            reason = line.split(":", 1)[1].strip()

    return AttributionResult(agent_name=agent_name, step_number=step_number, reason=reason)


# ---------------------------------------------------------------------------
# Method 1: All-at-Once
# ---------------------------------------------------------------------------

def all_at_once(
    trace: dict[str, Any],
    model: str = DEFAULT_MODEL,
    max_tokens: int = DEFAULT_MAX_TOKENS,
) -> AttributionResult:
    """Analyze a trace by feeding the entire conversation to the LLM at once.

    Args:
        trace: dict with required key "history" (list of conversation steps).
               Optional keys: "question", "ground_truth".
        model: LLM model identifier (any litellm-supported string).
        max_tokens: max response tokens.

    Returns:
        AttributionResult with agent_name, step_number, reason.
    """
    history = trace["history"]
    agent_key = _detect_agent_key(history)
    problem = trace.get("question", "")
    ground_truth = trace.get("ground_truth", "")

    chat_content = "\n".join(
        f"{entry.get(agent_key, 'Unknown Agent')}: {entry.get('content', '')}"
        for entry in history
    )

    prompt = (
        "You are an AI assistant tasked with analyzing a multi-agent conversation history when solving a real world problem. "
        f"The problem is: {problem}\n"
    )
    if ground_truth:
        prompt += f"The Answer for the problem is: {ground_truth}\n"
    prompt += (
        "Identify which agent made an error, at which step, and explain the reason for the error. "
        "Here's the conversation:\n\n" + chat_content +
        "\n\nBased on this conversation, please predict the following:\n"
        "1. The name of the agent who made a mistake that should be directly responsible for the wrong solution to the real world problem. "
        "If there are no agents that make obvious mistakes, decide one single agent in your mind. Directly output the name of the Expert.\n"
        "2. In which step the mistake agent first made mistake. For example, in a conversation structured as follows: "
        '{"agent a": "xx", "agent b": "xxxx", "agent c": "xxxxx", "agent a": "xxxxxxx"}, '
        "each entry represents a 'step' where an agent provides input. If the mistake is in agent c's speech, the step number is 2. "
        "If the second speech by 'agent a' contains the mistake, the step number is 3, and so on. "
        "Please determine the step number where the first mistake occurred.\n"
        "3. The reason for your prediction.\n"
        "Please answer in the format:\nAgent Name: (Your prediction)\nStep Number: (Your prediction)\nReason for Mistake: (Your explanation)\n"
    )

    messages = [
        {"role": "system", "content": "You are a helpful assistant skilled in analyzing conversations."},
        {"role": "user", "content": prompt},
    ]

    response = chat(messages, model=model, max_tokens=max_tokens)
    return _parse_all_at_once_result(response, history, agent_key)


# ---------------------------------------------------------------------------
# Method 2: Step-by-Step
# ---------------------------------------------------------------------------

def step_by_step(
    trace: dict[str, Any],
    model: str = DEFAULT_MODEL,
    max_tokens: int = DEFAULT_MAX_TOKENS,
) -> AttributionResult:
    """Analyze a trace step by step, asking the LLM at each step if an error occurred.

    Returns as soon as the first error is detected.
    """
    history = trace["history"]
    agent_key = _detect_agent_key(history)
    problem = trace.get("question", "")
    ground_truth = trace.get("ground_truth", "")

    current_conversation = ""

    for idx, entry in enumerate(history):
        agent_name = entry.get(agent_key, "Unknown Agent")
        content = entry.get("content", "")
        current_conversation += f"Step {idx} - {agent_name}: {content}\n"

        prompt = (
            f"You are an AI assistant tasked with evaluating the correctness of each step in an ongoing "
            f"multi-agent conversation aimed at solving a real-world problem. The problem being addressed is: {problem}. "
        )
        if ground_truth:
            prompt += f"The Answer for the problem is: {ground_truth}\n"
        prompt += (
            f"Here is the conversation history up to the current step:\n{current_conversation}\n"
            f"The most recent step ({idx}) was by '{agent_name}'.\n"
            f"Your task is to determine whether this most recent agent's action (Step {idx}) contains an error "
            "that could hinder the problem-solving process or lead to an incorrect solution. "
            "Please respond with 'Yes' or 'No' and provide a clear explanation for your judgment. "
            "Note: Please avoid being overly critical in your evaluation. Focus on errors that clearly derail the process.\n"
            "Respond ONLY in the format: 1. Yes/No.\n2. Reason: [Your explanation here]"
        )

        messages = [
            {"role": "system", "content": "You are a precise step-by-step conversation evaluator."},
            {"role": "user", "content": prompt},
        ]

        answer = chat(messages, model=model, max_tokens=max_tokens)

        if answer.lower().strip().startswith("1. yes"):
            reason_parts = answer.split("Reason:", 1)
            reason = reason_parts[-1].strip() if len(reason_parts) > 1 else answer
            return AttributionResult(agent_name=agent_name, step_number=idx, reason=reason)

    # No error found — report the last step as most suspicious
    last = history[-1]
    return AttributionResult(
        agent_name=last.get(agent_key, "Unknown Agent"),
        step_number=len(history) - 1,
        reason="No decisive error detected by step-by-step analysis; last step reported as most suspicious.",
    )


# ---------------------------------------------------------------------------
# Method 3: Binary Search
# ---------------------------------------------------------------------------

def binary_search(
    trace: dict[str, Any],
    model: str = DEFAULT_MODEL,
    max_tokens: int = DEFAULT_MAX_TOKENS,
) -> AttributionResult:
    """Analyze a trace using binary search to locate the error step efficiently."""
    history = trace["history"]
    agent_key = _detect_agent_key(history)
    problem = trace.get("question", "")
    ground_truth = trace.get("ground_truth", "")

    step = _binary_search_recursive(
        history=history,
        agent_key=agent_key,
        problem=problem,
        ground_truth=ground_truth,
        start=0,
        end=len(history) - 1,
        model=model,
        max_tokens=max_tokens,
    )

    entry = history[step]
    return AttributionResult(
        agent_name=entry.get(agent_key, "Unknown Agent"),
        step_number=step,
        reason=f"Binary search identified step {step} as the decisive error step.",
    )


def _binary_search_recursive(
    history: list[dict],
    agent_key: str,
    problem: str,
    ground_truth: str,
    start: int,
    end: int,
    model: str,
    max_tokens: int,
) -> int:
    """Recursively narrow down the error location via binary search."""
    if start >= end:
        return start

    segment = history[start: end + 1]
    chat_content = "\n".join(
        f"{entry.get(agent_key, 'Unknown Agent')}: {entry.get('content', '')}"
        for entry in segment
    )

    mid = start + (end - start) // 2
    range_desc = f"from step {start} to step {end}"
    upper_desc = f"from step {start} to step {mid}"
    lower_desc = f"from step {mid + 1} to step {end}"

    prompt = (
        "You are an AI assistant tasked with analyzing a segment of a multi-agent conversation. "
        "Multiple agents are collaborating to address a user query, with the goal of resolving the query through their collective dialogue.\n"
        "Your primary task is to identify the location of the most critical mistake within the provided segment. "
        "Determine which half of the segment contains the single step where this crucial error occurs, "
        "ultimately leading to the failure in resolving the user's query.\n"
        f"The problem to address is as follows: {problem}\n"
    )
    if ground_truth:
        prompt += f"The Answer for the problem is: {ground_truth}\n"
    prompt += (
        f"Review the following conversation segment {range_desc}:\n\n{chat_content}\n\n"
        f"Based on your analysis, predict whether the most critical error is more likely to be located "
        f"in the upper half ({upper_desc}) or the lower half ({lower_desc}) of this segment.\n"
        "Please provide your prediction by responding with ONLY 'upper half' or 'lower half'. "
        "Remember, your answer should be based on identifying the mistake that directly contributes to the failure. "
        "If no single clear error is evident, consider the step you believe is most responsible for the failure."
    )

    messages = [
        {"role": "system", "content": "You are an AI assistant specializing in localizing errors in conversation segments."},
        {"role": "user", "content": prompt},
    ]

    result = chat(messages, model=model, max_tokens=max_tokens)
    result_lower = result.lower()

    if "upper half" in result_lower:
        return _binary_search_recursive(history, agent_key, problem, ground_truth, start, mid, model, max_tokens)
    elif "lower half" in result_lower:
        return _binary_search_recursive(history, agent_key, problem, ground_truth, min(mid + 1, end), end, model, max_tokens)
    else:
        # Ambiguous — random fallback
        if random.randint(0, 1) == 0:
            return _binary_search_recursive(history, agent_key, problem, ground_truth, start, mid, model, max_tokens)
        else:
            return _binary_search_recursive(history, agent_key, problem, ground_truth, min(mid + 1, end), end, model, max_tokens)
