import concurrent.futures
import logging
from typing import Any

from tqdm import tqdm

from deception_detection.data.base import DialogueDataset
from deception_detection.log import logger
from deception_detection.rollouts.api import grade_single_dialogue_for_deception
from deception_detection.scores import Scores
from deception_detection.types import Dialogue, Judgement, Message
from deception_detection.utils import REPO_ROOT

BLACK_BOX_PROMPT_FILE = REPO_ROOT / "data" / "black_box_baseline" / "prompt.txt"
BLACK_BOX_PROMPT = BLACK_BOX_PROMPT_FILE.read_text().strip()


def merge_messages(dialogue: list[Message]) -> list[Message]:
    """Merge consecutive messages from the same role"""
    merged_dialogue = []
    for message in dialogue:
        if not merged_dialogue or merged_dialogue[-1].role != message.role:
            merged_dialogue.append(message)
        else:
            merged_dialogue[-1].content += "\n" + message.content
    return merged_dialogue


def format_dialogue(dialogue: list[Message]) -> str:
    """Format a dialogue for the black box prompt"""
    merged_dialogue = merge_messages(dialogue)
    return "\n".join(f"{m.role}: {m.content}" for m in merged_dialogue)


def black_box_grading_prompt_getter(
    input_dialogue: Dialogue, response: Message, metadata: dict[str, Any] | None
) -> dict[str, str]:
    """Get the grading prompt for a single dialogue"""
    out = {
        "system": BLACK_BOX_PROMPT,
        "user": f"<context>\n{format_dialogue(input_dialogue)}\n</context>\n"
        f"\n<output>\n{response.content}\n</output>",
    }
    return out


def compute_black_box_scores_for_dataset(
    dataset: DialogueDataset, mock: bool = False
) -> tuple[Scores, list[str]]:
    """
    Compute black box scores for a single dataset by grading each dialogue with an API model.
    """
    logger.info(f"Computing black box scores for dataset {dataset.name}...")
    # suppress httpx info-level logs from anthropic API calls
    logging.getLogger("httpx").setLevel(logging.WARNING)

    def get_context(dialogue: list[Message]) -> list[Message]:
        """Get all messages except the final assistant response(s)"""
        # Find the last non-assistant message
        last_non_assistant_idx = len(dialogue) - 1
        while last_non_assistant_idx >= 0 and dialogue[last_non_assistant_idx].role == "assistant":
            last_non_assistant_idx -= 1
        return dialogue[: last_non_assistant_idx + 1]

    def get_assistant_response(dialogue: list[Message]) -> Message:
        """Get the final assistant response(s) glued together"""
        # Find the last non-assistant message
        last_non_assistant_idx = len(dialogue) - 1
        while last_non_assistant_idx >= 0 and dialogue[last_non_assistant_idx].role == "assistant":
            last_non_assistant_idx -= 1

        # Combine all subsequent assistant messages
        combined_content = "".join(m.content for m in dialogue[last_non_assistant_idx + 1 :])
        return Message(role="assistant", content=combined_content, detect=True)

    judgements: list[Judgement] = []
    reasonings: list[str] = []
    futures_and_indices = []

    with concurrent.futures.ThreadPoolExecutor(max_workers=50) as executor:
        # Create futures with their original indices
        for i, dialogue in enumerate(dataset.dialogues):
            future = executor.submit(
                grade_single_dialogue_for_deception,
                black_box_grading_prompt_getter,
                get_context(dialogue),
                get_assistant_response(dialogue),
                None,
                verbose=False,
                mock=mock,
                api="openai",
            )
            futures_and_indices.append((future, i))

        # Initialize judgements list with None placeholders
        judgements = [None] * len(dataset.dialogues)  # type: ignore
        reasonings = [None] * len(dataset.dialogues)  # type: ignore

        # Add progress bar
        for future, idx in tqdm(
            [(f, i) for f, i in futures_and_indices],  # type: ignore
            total=len(futures_and_indices),  # type: ignore
            desc="Grading dialogues",
        ):
            judgement, reasoning = future.result()
            judgements[idx] = judgement
            reasonings[idx] = reasoning

    def to_score(j: Judgement) -> float:
        if j == "deceptive":
            return 1.0
        elif j == "honest":
            return 0.0
        else:
            try:
                return float(j)
            except ValueError:
                return float("nan")

    scores = [[to_score(j)] for j in judgements]
    return Scores(scores, dataset.labels), reasonings
