import csv
import os
from pathlib import Path
from random import shuffle
from typing import Literal

import replicate
from groq import Groq
from openai import OpenAI
from tqdm import tqdm


def make_adjectives_list():
    adjective_path = Path("../data/adjectives.csv")

    with open(adjective_path, 'r') as file:
        reader = csv.reader(file)
        adjectives = [row[0] for row in reader]

    noun_path = Path("../data/nouns.csv")
    with open(noun_path, 'r') as file:
        reader = csv.reader(file)
        nouns = [row[0] for row in reader]


    shuffle(adjectives)
    shuffle(nouns)

    pairs = [f"{adjectives[idx]} {nouns[idx]}" for idx in range(len(adjectives))]

    with open("../data/adjective_noun_pairs.txt", "w") as fp:
        fp.write("\n".join(pairs))


def filter_bad_and_nonimageable_words():
    with open("../data/imageable_deduplicate_clean_v1.txt", "r") as fp:
        adj_noun_pairs = fp.read().split("\n")
    with open("../data/discarded_words.txt", "r") as fp:
        discarded_words = fp.read().split("\n")

    output_file = Path("../data/imageable_deduplicate_clean_v2.txt")
    existing_pairs = set()
    if output_file.exists():
        with open(output_file, 'r') as fp:
            existing_pairs = set(fp.read().split("\n"))

    with open(output_file, "a") as fp:
        for pair in tqdm(adj_noun_pairs):
            if pair in existing_pairs:
                continue
            valid = valid_noun_openai(pair)
            #valid = valid_noun_llama3(pair)
            if valid == "false" or valid == "maybe":
                print(f"{valid}: {pair}")
                continue
            fp.write(f"{pair}\n")


def valid_noun_openai(noun: str) -> Literal["true", "false", "maybe"]:
    client = OpenAI(
        api_key=os.getenv("XAI_API_KEY"),
        base_url = "https://api.x.ai/v1",
    )
    prompt = f"Is this word an object that can be made into an image and is the word appropriate for children? Your response should be 'true' or 'false'. The Word: {noun}"
    completion = client.chat.completions.create(
        model="grok-2-1212",
        messages=[
            {"role": "assistant", "content": "You are a helpful assistant."},
            {"role": "user", "content": prompt}
        ]
    )
    response = completion.choices[0].message.content

    if "false" in response.lower():
        return "false"
    elif "true" in response.lower():
        return "true"

    return "maybe"

def valid_noun_llama3(noun: str) -> Literal["true", "false", "maybe"]:
    prompt = f"Is this word an object that can be made into an image? Is the word appropriate for children? Your response should be 'true' or 'false'. The Word: {noun}"
    response = ""
    # The meta/meta-llama-3-8b model can stream output as it's running.
    for event in replicate.stream(
            "meta/meta-llama-3-70b-instruct",
            input={
                "top_k": 0,
                "top_p": 0.95,
                "prompt": prompt,
                "max_tokens": 512,
                "temperature": 0.7,
                "system_prompt": "You are a helpful assistant.",
                "length_penalty": 1,
                "max_new_tokens": 512,
                "stop_sequences": "<|end_of_text|>,<|eot_id|>",
                "prompt_template": "<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\n{system_prompt}<|eot_id|><|start_header_id|>user<|end_header_id|>\n\n{prompt}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n",
                "presence_penalty": 0,
                "log_performance_metrics": False
            },
    ):
        response += str(event)

    if "true" in response.lower():
        return "true"
    elif "false" in response.lower():
        return "false"

    return "maybe"


def filter_invalid_nouns_openai(nouns: list[str]) -> list[str]:
    client = OpenAI(
        api_key=os.getenv("XAI_API_KEY"),
        base_url = "https://api.x.ai/v1",
    )
    prompt = f"Below is a list of words. Your job is to list any words that are inappropriate for children in particular (sex, drugs, alcohol, and violence). Your response should should be a newline delimited list of words. Please don't add any extra details.\n\n{"\n".join(nouns)}"
    completion = client.chat.completions.create(
        model="grok-2-1212",
        messages=[
            {"role": "assistant", "content": "You are a helpful assistant."},
            {"role": "user", "content": prompt}
        ]
    )
    response = completion.choices[0].message.content
    invalid_words = []
    for noun in nouns:
        if noun in response:
            invalid_words.append(noun)

    return invalid_words


if __name__ == "__main__":
    with open("../data/imageable_deduplicate.txt", "r") as fp:
        adj_noun_pairs = fp.read().split("\n")

    with open("../data/bad_nouns_v2.txt", "r") as fp:
        bad_words = fp.read().split("\n")
    adj_noun_pairs = list(set(adj_noun_pairs) - set(bad_words))

    bad_nouns = []
    subset_len = 20
    for start in tqdm(range(0, len(adj_noun_pairs), subset_len)):
        end = min(len(adj_noun_pairs), start+subset_len)
        subset = adj_noun_pairs[start:end]
        bad_nouns.extend(filter_invalid_nouns_openai(subset))

    print(f"adding {len(bad_nouns)} new words")

    with open("../data/bad_nouns_v2.txt", "a") as fp:
        fp.write("\n".join(bad_nouns))
