import asyncio
import json
import os
import sys
from os import getenv
from pathlib import Path

import evaluate
import pandas as pd
import requests
from aiolimiter import AsyncLimiter
from dotenv import load_dotenv
from joblib.memory import Memory
from openai import AsyncOpenAI
from tqdm.asyncio import tqdm_asyncio
from transformers import NllbTokenizer

# config
models = [
    "openai/gpt-4o-mini",
    "anthropic/claude-3.5-haiku",
    # "meta-llama/llama-3.1-405b-instruct",  # lots of slow repetitions for LRLs
    # "mistralai/mistral-large",
    # "google/gemini-flash-1.5",  # very fast
    # "qwen/qwen-2.5-72b-instruct",  # somewhat slow
]
fast_model = "anthropic/claude-3.5-haiku"
n_sentences = 30

# setup
load_dotenv()
client = AsyncOpenAI(
    base_url="https://openrouter.ai/api/v1",
    api_key=getenv("OPENROUTER_API_KEY"),
)
cache = Memory(location=".cache", verbose=0).cache
bleu = evaluate.load("bleu")
bertscore = evaluate.load("bertscore")
tokenizer = NllbTokenizer.from_pretrained("facebook/nllb-200-distilled-600M")
rate_limit = AsyncLimiter(max_rate=20, time_period=1)


def reorder(language_name):
    if "," in language_name and "(" not in language_name:
        return language_name.split(",")[1] + " " + language_name.split(",")[0]
    return language_name

# load benchmark languages and scripts
data = Path("src/data")
benchmark_dir = data / "floresp-v2.0-rc.3/dev"
benchmark_languages = pd.DataFrame(
    [f.split(".")[1].split("_", 1) for f in os.listdir(benchmark_dir)],
    columns=["language_code", "script_code"],
)
# hack: drop additional script codes for languages with multiple scripts
benchmark_languages = benchmark_languages.groupby("language_code").head(1)
benchmark_languages["in_benchmark"] = True

# load Ethnologue language names
language_names = (
    pd.read_csv(data / "LanguageCodes.tab", sep="\t")
    .rename(columns={"LangID": "language_code", "Name": "language_name"})[
        ["language_code", "language_name"]
    ]
    .assign(language_name=lambda df: df["language_name"].apply(reorder).str.strip())
)

# load Wikidata speaker stats
language_stats = (
    pd.read_csv(data / "languages.tsv", sep="\t")
    .rename(columns={"iso639_3": "language_code", "maxSpeakers": "speakers"})[
        ["language_code", "speakers"]
    ]
    .dropna(subset=["language_code"])
)
language_stats["speakers"] = pd.to_numeric(language_stats["speakers"], errors="coerce")
ignored_languages = [
    "zho",  # Chinese -> use Mandarin (cmn) instead
    "ara",  # Arabic -> use Standard Arabic (arb) instead
    "pus",  # Pashto -> use Nothern / Central / Southern Pashto instead (pbt / pst / pbu)
    "fas",  # Persian -> use Iranian Persian (pes) instead
    "msa",  # Malay -> use Indonesian (ind) instead
]
language_stats = language_stats[
    ~language_stats["language_code"].isin(ignored_languages)
]

# load unicode script names
script_names = pd.read_csv(data / "ScriptCodes.csv").rename(
    columns={"Code": "script_code", "English Name": "script_name"}
)[["script_code", "script_name"]]

# merge data
languages = pd.merge(language_stats, language_names, on="language_code", how="outer")
languages = pd.merge(benchmark_languages, languages, on="language_code", how="outer")
languages = pd.merge(languages, script_names, on="script_code", how="left")
languages["in_benchmark"] = languages["in_benchmark"].fillna(False)
languages = languages.sort_values(by="speakers", ascending=False)

# sample languages to translate from
# when translating e.g. to Mandarin, we drop Mandarin from the sample and use the next samples from the list instead; therefore we need to sample more than n_sentences
original_languages = languages[languages["in_benchmark"]].sample(
    n=n_sentences * 2, weights="speakers", replace=True, random_state=42
)
# sample languages to analyze with all models
detailed_target_languages = languages[languages["in_benchmark"]].sample(
    n=3, random_state=42
)


# utils
def check_rate_limit():
    print(
        requests.get(
            "https://openrouter.ai/api/v1/auth/key",
            headers={"Authorization": f"Bearer {getenv('OPENROUTER_API_KEY')}"},
        ).json()
    )
    models = requests.get(
        "https://openrouter.ai/api/v1/models",
        headers={"Authorization": f"Bearer {getenv('OPENROUTER_API_KEY')}"},
    ).json()["data"]
    model = next((m for m in models if m["id"] == "google/gemini-flash-1.5"), None)
    print(model)


@cache
async def complete(**kwargs):
    async with rate_limit:
        response = await client.chat.completions.create(**kwargs)
    if not response.choices:
        raise Exception(response)
    return response


@cache
async def translate(model, target_language, target_script, sentence):
    reply = await complete(
        model=model,
        messages=[
            {
                "role": "user",
                "content": f"Translate the following text to the {target_language} language; use the {target_script} script; reply only with the translation:\n\n{sentence}",
            }
        ],
        temperature=0,
        max_tokens=1024,
    )
    return reply.choices[0].message.content


def mean(l):
    return sum(l) / len(l) if l else 0


def load_sentences(language):
    return open(
        f"{benchmark_dir}/dev.{language.language_code}_{language.script_code}"
    ).readlines()


# evaluation!
async def main():
    results = []
    for language in list(languages.itertuples())[:5]:
        name = (
            language.language_name
            if not pd.isna(language.language_name)
            else language.language_code
        )
        print(name, file=sys.stderr)
        scores = []
        if language.in_benchmark:
            target_sentences = load_sentences(language)[:n_sentences]
            for model in models:
                if (
                    model != fast_model
                    and language.language_code
                    not in detailed_target_languages.language_code.values
                ):
                    continue
                # drop the target language from the original languages sample
                _original_languages = original_languages[
                    original_languages.language_code != language.language_code
                ].iloc[:n_sentences]
                original_sentences = [
                    load_sentences(lang)[i]
                    for i, lang in enumerate(_original_languages.itertuples())
                ]
                print(model, file=sys.stderr)
                predictions = [
                    translate(
                        model, language.language_name, language.script_name, sentence
                    )
                    for sentence in original_sentences
                ]
                predictions = await tqdm_asyncio.gather(*predictions, miniters=1)
                metrics_bleu = bleu.compute(
                    predictions=predictions,
                    references=target_sentences,
                    tokenizer=tokenizer.tokenize,
                )
                # metrics_bert = bertscore.compute(
                #     predictions=predictions,
                #     references=target_sentences,
                #     model_type="distilbert-base-uncased",
                # )
                scores.append(
                    {
                        "model": model,
                        "bleu": metrics_bleu["bleu"],
                        # "bert_score": mean(metrics_bert["f1"]),
                    }
                )
        results.append(
            {
                "language_name": name,
                "language_code": language.language_code,
                "speakers": language.speakers if not pd.isna(language.speakers) else 0,
                "scores": scores,
                "bleu": mean([s["bleu"] for s in scores]) or -0.02,
                # "bert_score": mean([s["bert_score"] for s in scores]),
            }
        )
    print(json.dumps(results, indent=2, ensure_ascii=False))


if __name__ == "__main__":
    # check_rate_limit()
    asyncio.run(main())