import gc

import weave
from hazm import sent_tokenize, word_tokenize as persian_word_tokenize
from nltk.tokenize import word_tokenize as english_word_tokenize
import PyPDF2
import numpy as np
import re
import wandb
from huggingface_hub import login
from collections import OrderedDict
import MetaRagTool.Utils.Constants as Constants
from MetaRagTool.Encoders.SentenceTransformerEncoder import SentenceTransformerEncoder
from MetaRagTool.LLM.GoogleGemini import Gemini
from MetaRagTool.LLM.JudgeLLM import JudgeLLM
from MetaRagTool.Utils import DataLoader
from MetaRagTool.Utils.MetaRagConfig import MetaRagConfig

has_logged_in_to_wandb = False


def init_wandb(project_name, run_name, config=None):
    global has_logged_in_to_wandb
    if not has_logged_in_to_wandb:
        api_key = Constants.WandbToken
        wandb.login(key=api_key)
        has_logged_in_to_wandb = True

    wandb.init(project=project_name, name=run_name, config=config)


def init_hf():
    hf_token = Constants.HFToken
    login(token=hf_token, add_to_git_credential=False)


def listToString(listOfStrings, separator="\n"):
    output = ""
    for s in listOfStrings:
        output += str(s) + separator
    return output


def capped_sent_tokenize(text, max_length=500, sentence_tokenizer=sent_tokenize,
                         custom_word_tokenizer=persian_word_tokenize):
    sentences = sentence_tokenizer(text)
    capped_sentences = []
    for sentence in sentences:
        while token_len(sentence) > max_length:
            tokens = custom_word_tokenizer(sentence)
            capped_sentences.append(" ".join(tokens[:max_length]))
            sentence = " ".join(tokens[max_length:])
        capped_sentences.append(sentence)
    return capped_sentences


def token_len(text, custom_word_tokenizer=persian_word_tokenize):
    return len(custom_word_tokenizer(text))


def token_len_en(text):
    return len(english_word_tokenize(text))


def reflect_vector(q, v):
    proj = np.dot(v, q) / np.dot(q, q) * q
    reflection = 2 * proj - v

    return reflection


def read_pdf(pdf_path, ignore_line_breaks=True):
    text = ""
    with open(pdf_path, 'rb') as file:
        pdf_reader = PyPDF2.PdfReader(file)
        for page in pdf_reader.pages:
            if page.extract_text():
                page_text = page.extract_text()
                if ignore_line_breaks:
                    # Replace single line breaks with space
                    page_text = re.sub(r'(?<!\n)\n(?!\n)', ' ', page_text)
                    text += page_text + '\n'
                else:
                    text += page_text

    return text.strip()


def Init(encoder_name, top_k, sample_size, qa_sample_ratio, local_mode=False, multi_hop_hardness_factor=0,
         judged=False, useTool=False, llm_name:str=Gemini.GEMINI_2_FLASH,multi_hop=True) -> MetaRagConfig:
    gc.collect()

    Constants.local_mode = local_mode
    Constants.use_wandb = True
    contexts, qas = DataLoader.loadWikiFaQa(sample_size=sample_size, qa_sample_ratio=qa_sample_ratio, multi_hop=multi_hop,
                                            multi_hop_hardness_factor=multi_hop_hardness_factor)
    encoder = SentenceTransformerEncoder(encoder_name)
    llm = None
    judge = None

    # wandb_config = {
    #     'model_name': llm_name,
    #     'multi_hop': True,
    #     'sample_size': sample_size,
    #     'qa_sample_ratio': qa_sample_ratio,
    #     'hardness_factor': multi_hop_hardness_factor,
    #     'top_k': top_k,
    #     'useTool': useTool
    # }

    if judged:
        llm = Gemini(has_memory=False, RequestPerMinute_limit=15, model_name=llm_name)
        judge = JudgeLLM(model_name=Gemini.GEMINI_2_FLASH_LIGHT, RequestPerMinute_limit=30)
        project_name = "fullRag"
        weave.init(project_name)

    else:
        project_name = 'retrival'

    ragConfig = MetaRagConfig(encoder_name=encoder_name, top_k=top_k, sample_size=sample_size,
                              qa_sample_ratio=qa_sample_ratio, local_mode=local_mode,
                              multi_hop_hardness_factor=multi_hop_hardness_factor, judged=judged, useTool=useTool,
                              llm_name=llm_name,
                              encoder=encoder, contexts=contexts, qas=qas, llm=llm, judge=judge,
                              project_name=project_name,multi_hop=multi_hop)

    return ragConfig


def remove_duplicates(l):
    return list(OrderedDict.fromkeys(l))
