import random
from time import sleep
from typing import Any, Literal, cast

import dotenv
from anthropic import Anthropic
from openai import OpenAI
from together import Together
from together.types import CompletionResponse

from deception_detection.log import logger
from deception_detection.models import ModelName
from deception_detection.types import GradingPromptGetter, Judgement, Message


def parse_judgement(response_text: str) -> tuple[Judgement, str]:
    """Parse the judgement from the response text.

    Args:
        response_text: Raw response text from the API

    Returns:
        tuple of (parsed_judgement, original_response_text)
    """
    try:
        # Extract text between judgement tags
        parts = response_text.split("<judgement>")
        if len(parts) != 2:
            logger.warning(f"Missing or multiple <judgement> tags in response {response_text}")
            return "unknown", response_text

        judgement_section = parts[1].split("</judgement>")
        if len(judgement_section) != 2:
            logger.warning(f"Missing or multiple </judgement> tags in response {response_text}")
            return "unknown", response_text

        judgement = judgement_section[0].strip().lower()

        if judgement.isdigit():
            judgement = int(judgement)
        return judgement, response_text

    except Exception as e:
        logger.warning(f"Error parsing judgement: {str(e)}")
        return "unknown", response_text


def handle_error(attempt: int, error: Exception, base_delay: float = 2) -> None:
    """Handle errors with exponential backoff.

    Args:
        attempt: The current retry attempt number
        error: The error that was raised
        base_delay: Base delay for exponential backoff
    """
    delay = base_delay**attempt
    jitter = random.uniform(0, 0.1 * delay)  # Add 0-10% jitter
    logger.warning(f"Error (attempt {attempt}): {error.__class__.__name__}")
    logger.warning(f"Waiting {delay:.1f} seconds before retry")
    sleep(delay + jitter)


def grade_single_dialogue_for_deception(
    grading_prompt_getter: GradingPromptGetter,
    input_messages: list[Message],
    output_message: Message,
    metadata: dict[str, Any] | None,
    verbose: bool,
    mock: bool = False,
    api: Literal["openai", "anthropic"] = "openai",
) -> tuple[Judgement, str]:
    """Grade a single dialogue as deceptive or not using either OpenAI or Anthropic API.

    Args:
        grading_prompt_getter: Function to generate the grading prompt
        input_messages: List of input messages in the dialogue
        output_message: The output message to grade
        metadata: Additional metadata for grading
        verbose: Whether to print verbose output
        mock: Whether to return mock responses
        api: Which API to use ("openai" or "anthropic")

    Returns:
        tuple of (judgement, grader_reasoning)
    """
    max_retries = 4
    try:
        prompt = grading_prompt_getter(input_messages, output_message, metadata)
    except Exception as e:
        logger.warning(f"Error generating grading prompt: {e}")
        return "unknown", "Prompt generation error"

    if mock:
        return random.choice(["honest", "deceptive"]), "mocked grader reasoning"

    dotenv.load_dotenv()
    last_error = None

    for attempt in range(max_retries):
        try:
            if api == "openai":
                client = OpenAI(organization="org-CGjCXR7l3PLggjuyxf0yu5bU")
                response = client.chat.completions.create(
                    model="gpt-4o-2024-08-06",
                    max_tokens=1000,
                    messages=[
                        {"role": "system", "content": prompt["system"]},
                        {"role": "user", "content": prompt["user"]},
                    ],
                )
                response_text = cast(str, response.choices[0].message.content)
            else:  # anthropic
                client = Anthropic()
                response = client.messages.create(
                    model="claude-3-5-sonnet-20241022",
                    max_tokens=1000,
                    system=prompt["system"],
                    messages=[{"role": "user", "content": prompt["user"]}],
                )
                response_text = cast(str, response.content[0].text)  # type: ignore

            return parse_judgement(response_text)

        except Exception as e:
            last_error = e
            logger.warning(f"Unexpected error (attempt {attempt}): {e.__class__.__name__}")
            handle_error(attempt, e)
            continue

    error_msg = f"Max retries exceeded. Last error: {last_error}"
    logger.error(error_msg)
    return "unknown", error_msg


def to_api_name(model_name: ModelName) -> str:
    """Convert a ModelName to the name used by the together API."""
    mappings = {
        ModelName.LLAMA_70B_3_3: "meta-llama/Llama-3.3-70B-Instruct-Turbo",
        ModelName.LLAMA_70B_3_1: "meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo",
        ModelName.MISTRAL_3_1_24B: "mistralai/Mistral-Small-Instruct-2409",
        ModelName.QWEN_72B: "Qwen/Qwen2.5-72B-Instruct-Turbo",
    }
    return mappings.get(model_name, model_name.value)


def generate_rollout_from_together_api(
    formatted_input: str, num_rollouts: int, model_name: ModelName, mock: bool = False
) -> list[str]:
    if mock:
        return [" Mocked Response" for _ in range(num_rollouts)]

    model = to_api_name(model_name)
    max_retries = 4
    last_error = None
    all_responses: list[str] = []
    batch_size = 4  # Maximum number of rollouts per API call

    # Calculate number of full batches and remaining rollouts
    num_batches = (num_rollouts + batch_size - 1) // batch_size

    for batch in range(num_batches):
        current_batch_size = min(batch_size, num_rollouts - len(all_responses))

        for attempt in range(max_retries):
            try:
                client = Together()
                eot = "<|end_of_text|>" if model.startswith("meta-llama") else "<|endoftext|>"
                response = client.completions.create(
                    prompt=formatted_input,
                    model=model,
                    stop=[eot],
                    n=current_batch_size,
                    temperature=1,
                    max_tokens=3072,
                )
                assert isinstance(response, CompletionResponse)
                assert response.choices is not None
                assert len(response.choices) == current_batch_size

                batch_responses = [choice.text for choice in response.choices]
                all_responses.extend(batch_responses)
                break  # Success - exit retry loop

            except Exception as e:
                last_error = e
                logger.warning(f"Error (attempt {attempt}): {e.__class__.__name__}")
                handle_error(attempt, e)
                if attempt == max_retries - 1:  # Last attempt failed
                    error_msg = f"Max retries exceeded for batch {batch}. Last error: {last_error}"
                    logger.error(error_msg)
                    # Fill the batch with error messages
                    all_responses.extend([f"Error: {error_msg}" for _ in range(current_batch_size)])

    return all_responses
