# %%
import random
from concurrent.futures import Future, ThreadPoolExecutor, as_completed
from pathlib import Path
from typing import cast

import pandas as pd
from anthropic import Anthropic
from tqdm import tqdm

from deception_detection.utils import REPO_ROOT

PATHS = [
    # REPO_ROOT / "data/geometry_of_truth/cities.csv",
    # REPO_ROOT / "data/geometry_of_truth/cities_cities_conj.csv",
    # REPO_ROOT / "data/geometry_of_truth/cities_cities_disj.csv",
    # REPO_ROOT / "data/geometry_of_truth/common_claim_true_false.csv",
    # REPO_ROOT / "data/geometry_of_truth/companies_true_false.csv",
    # REPO_ROOT / "data/geometry_of_truth/counterfact_true_false.csv",
    REPO_ROOT / "data/geometry_of_truth/larger_than.csv",
    REPO_ROOT / "data/geometry_of_truth/neg_cities.csv",
    REPO_ROOT / "data/geometry_of_truth/neg_sp_en_trans.csv",
    REPO_ROOT / "data/geometry_of_truth/smaller_than.csv",
    REPO_ROOT / "data/geometry_of_truth/sp_en_trans.csv",
    REPO_ROOT / "data/internal_state/animals_true_false.csv",
    REPO_ROOT / "data/internal_state/cities_true_false.csv",
    REPO_ROOT / "data/internal_state/companies_true_false.csv",
    REPO_ROOT / "data/internal_state/elements_true_false.csv",
    REPO_ROOT / "data/internal_state/facts_true_false.csv",
    REPO_ROOT / "data/internal_state/inventions_true_false.csv",
]

PROMPT = Path(REPO_ROOT / "scripts/conversion_prompt.txt").read_text()

# %%

# DO NOT RUN AGAIN


def shuffle_and_save(file_path: Path) -> None:
    # Read the CSV file
    df = pd.read_csv(file_path)

    # Shuffle the rows
    df = df.sample(frac=1, random_state=random.randint(1, 1000)).reset_index(drop=True)

    # Write the shuffled data back to the same file
    df.to_csv(file_path, index=False)

    # log
    print(f"Shuffled and saved: {file_path}")


# Shuffle and save each file
# for path in PATHS:
#     shuffle_and_save(path)
#     print(f"Shuffled and saved: {path}")

# %%


def process_row(client: Anthropic, statement: str, label: str) -> tuple[str, str, str, str, str]:
    """Process a single row using the Anthropic API."""
    prompt = PROMPT.format(statement=statement, label=label)

    response = client.messages.create(
        model="claude-3-haiku-20240307",
        max_tokens=300,
        system="You are a helpful assistant.",
        messages=[{"role": "user", "content": prompt}],
    )

    response_text = cast(str, response.content[0].text)  # type: ignore

    question = response_text.split("<question>")[1].split("</question>")[0].strip()
    correct_answer = (
        response_text.split("<correct_answer>")[1].split("</correct_answer>")[0].strip()
    )
    plausible_answer = (
        response_text.split("<plausible_answer>")[1].split("</plausible_answer>")[0].strip()
    )

    return (statement, label, question, correct_answer, plausible_answer)


def process_file_with_anthropic(
    file_path: Path, num_rows: int = 1000, max_workers: int = 10
) -> list[tuple[str, str, str, str, str]]:
    """
    Process the first num_rows of a CSV file using the Anthropic API with multithreading.

    Args:
    file_path (Path): Path to the CSV file.
    num_rows (int): Number of rows to process. Defaults to 1000.
    max_workers (int): Maximum number of threads to use. Defaults to 10.

    Returns:
    List[Tuple[str, str, str, str, str]]: List of tuples containing
    (statement, label, question, correct_answer, plausible_answer)
    """
    # Read the CSV file
    df = pd.read_csv(file_path)

    # Take the first num_rows
    df = df.head(num_rows)

    results = []
    client = Anthropic()

    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        future_to_row: dict[Future[tuple[str, str, str, str, str]], pd.Series] = {  # type: ignore
            executor.submit(
                process_row, client, cast(str, row["statement"]), cast(str, row["label"])
            ): row
            for _, row in df.iterrows()
        }
        for future in tqdm(
            as_completed(future_to_row),  # type: ignore
            total=len(future_to_row),  # type: ignore
            desc="Processing rows",
        ):
            try:
                result = future.result()
                results.append(result)
            except Exception as exc:
                print(f"Generated an exception: {exc}")

    return results


def save_results_to_csv(
    results: list[tuple[str, str, str, str, str]], original_file_path: Path
) -> None:
    """
    Save the processed results to a new CSV file with '_qa' appended to the original filename.

    Args:
    results (list[tuple[str, str, str, str, str]]): List of tuples containing
        (statement, label, question, correct_answer, plausible_answer)
    original_file_path (Path): Path to the original CSV file.

    Returns:
    None
    """
    # Create the new filename by appending '_qa' before the extension
    new_file_path = original_file_path.with_name(
        f"{original_file_path.stem}_qa{original_file_path.suffix}"
    )

    # Convert the results to a DataFrame
    df = pd.DataFrame(
        results,
        columns=[
            "statement",
            "label",
            "question",
            "correct_answer",
            "plausible_answer",
        ],  # type: ignore
    )

    # Save the DataFrame to the new CSV file
    df.to_csv(new_file_path, index=False)
    print(f"Results saved to {new_file_path}")


def run_conversion(paths: list[Path], num_rows: int = 1000, max_workers: int = 10) -> None:
    """
    Run the conversion process for each file path in the given list.

    Args:
    paths (list[Path]): List of file paths to process.
    num_rows (int): Number of rows to process for each file. Defaults to 1000.
    max_workers (int): Maximum number of threads to use. Defaults to 10.

    Returns:
    None
    """
    for path in paths:
        print(f"Processing file: {path}")
        results = process_file_with_anthropic(path, num_rows=num_rows, max_workers=max_workers)
        save_results_to_csv(results, path)
        print(f"Finished processing {path}\n")


run_conversion(PATHS, num_rows=1000, max_workers=10)


# %%
