import logging
import time
from decimal import Decimal

from tqdm.auto import tqdm

# RAGWithUsage lives in rag_usage (production code); re-exported here so that
# existing notebooks that do `from evaluation.evaluation_utils import RAGWithUsage`
# continue to work without changes.
from rag.rag_usage import RAGWithUsage  # noqa: F401  (re-export)
from config import MODEL_NAME

logger = logging.getLogger(__name__)


def calc_price(usage):
    input_price_per_million = 0.1
    output_price_per_million = 0.1

    input_cost = (usage.prompt_tokens / 1_000_000) * input_price_per_million
    output_cost = (usage.completion_tokens / 1_000_000) * output_price_per_million
    total_cost = input_cost + output_cost

    return {
        "input_cost": round(Decimal(input_cost), 7),
        "output_cost": round(Decimal(output_cost), 7),
        "total_cost": round(Decimal(total_cost), 7),
    }


def calc_total_price(usages):
    total_cost = 0.0

    for usage in usages:
        cost = calc_price(usage)
        total_cost = total_cost + float(cost["total_cost"])

    return total_cost


def llm_structured(
    client, instructions, user_prompt, output_type, model=MODEL_NAME
):

    messages = [
        {"role": "developer", "content": instructions},
        {"role": "user", "content": user_prompt},
    ]

    response = client.completion(
        model=model, messages=messages, response_format=output_type
    )

    output_parsed = output_type.model_validate_json(response.choices[0].message.content)

    return output_parsed, response.usage


def llm_structured_retry(
    client,
    instructions,
    user_prompt,
    output_type,
    model=MODEL_NAME,
    max_retries=3,
):
    for attempt in range(max_retries):
        try:
            return llm_structured(
                client,
                instructions,
                user_prompt,
                output_type,
                model=model,
            )
        except Exception as exc:
            if attempt == max_retries - 1:
                raise
            logger.warning(
                "llm_structured_retry attempt %d/%d failed: %s",
                attempt + 1,
                max_retries,
                exc,
            )
            time.sleep(2**attempt)


def map_progress(pool, seq, f, llm_client, data_gen_instructions, Questions):
    results = []

    with tqdm(total=len(seq)) as progress:
        futures = []

        for el in seq:
            future = pool.submit(f, el, llm_client, data_gen_instructions, Questions)
            future.add_done_callback(lambda p: progress.update())
            futures.append(future)

        for future in futures:
            result, usage = future.result()
            results.append(result)

    return results

def map_progress_answers(pool, seq, f, doc_idx, assistant):
    results = []

    with tqdm(total=len(seq)) as progress:
        futures = []

        for el in seq:
            future = pool.submit(f, el, doc_idx, assistant)
            future.add_done_callback(lambda p: progress.update())
            futures.append(future)

        for future in futures:
            result = future.result()
            results.append(result)

    return results

def map_progress_judge(
        pool,
        seq,
        f,
        litellm_client,
        aqa_judge_instructions,
        aqa_judge_prompt,
        AnswerEvaluation
    ):
    results = []
    usages = []

    with tqdm(total=len(seq)) as progress:
        futures = []

        for el in seq:
            future = pool.submit(
                f,
                el,
                litellm_client,
                aqa_judge_instructions,
                aqa_judge_prompt,
                AnswerEvaluation
            )
            future.add_done_callback(lambda p: progress.update())
            futures.append(future)

        for future in futures:
            result, usage = future.result()
            results.append(result)
            usages.append(usage)

    return results, usages