# llm.py (адаптировано под LangChain >= 1.0)
import os
import json
from mistralai.client import Mistral
from langchain_community.tools import TavilySearchResults, JinaSearch
from langchain_text_splitters import CharacterTextSplitter
from langchain_mistralai import ChatMistralAI
from langchain_core.prompts import PromptTemplate
from langchain_core.documents import Document
from transformers import AutoTokenizer
from utils.search_article import *
import concurrent.futures

tokenizer = AutoTokenizer.from_pretrained("mistral-community/pixtral-12b")

if not os.environ.get("TAVILY_API_KEY"):
    os.environ["TAVILY_API_KEY"] = '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'

api_key_2 = os.getenv("MISTRAL_API_KEY_2")
client_2 = Mistral(api_key=api_key_2)

api_key_3 = os.getenv("MISTRAL_API_KEY_3")
client_3 = Mistral(api_key=api_key_3)

api_key_4 = os.getenv("MISTRAL_API_KEY_4")
client_4 = ChatMistralAI(api_key=api_key_4, model="pixtral-12b-2409")


def count_tokens_in_text(text):
    tokens = tokenizer(text, return_tensors="pt", truncation=False, add_special_tokens=True)
    return len(tokens["input_ids"][0])


def setup_search(question):
    try:
        tavily_tool = TavilySearchResults(max_results=20)
        results = tavily_tool.invoke({"query": f"{question}"})
        if isinstance(results, list):
            return results, 'tavily_tool'
        else:
            print("Unexpected format from TavilySearchResults:", results)
    except Exception as e:
        print("Error with TavilySearchResults:", e)

    try:
        jina_tool = JinaSearch()
        results = json.loads(str(jina_tool.invoke({"query": f"{question}"})))
        if isinstance(results, list):
            return results, 'jina_tool'
        else:
            print("Unexpected format from JinaSearch:", results)
    except Exception as e:
        print("Error with JinaSearch:", e)

    return [], ''


def ask_question_to_mistral(text, question, context, images=[]):
    prompt = f"Answer the following question without mentioning it or repeating the original text on which the question is asked in style markdown. IN RUSSIAN:\nQuestion: {question}\n\nText:\n{text}"
    message_content = [{"type": "text", "text": prompt}] + images

    response = client_2.chat.complete(
        model="pixtral-12b-2409",
        messages=[{"role": "user", "content": f'{message_content}\n\nAdditional Context from Web Search:\n{context}'}],
    )
    return response.choices[0].message.content


def process_article_for_summary(text, preferences, images=[], compression_percentage=30):
    prompt = f"""
    You are a commentator.
    # article:
    {text}

    # Instructions:
    ## Summarize:
    In clear and concise language, summarize the key points and themes presented in the article by cutting it by {compression_percentage} percent in the markdown format.
    Taking into account the user's wishes "{preferences}"
    """
    message_content = [{"type": "text", "text": prompt}] + images
    response = client_3.chat.complete(
        model="pixtral-12b-2409",
        messages=[{"role": "user", "content": message_content}]
    ).choices[0].message.content
    return response


def process_large_article_for_summary(text, preferences, images=[], compression_percentage=30):
    text_splitter = CharacterTextSplitter.from_huggingface_tokenizer(
        tokenizer, chunk_size=100000, chunk_overlap=14000
    )
    docs = text_splitter.create_documents([text])

    map_template = f"""На основе приведенного текста, выполните сжатие, выделяя основные темы и важные моменты.
    Уровень сжатия: {compression_percentage}%.
    С учетом пожеланий пользователя "{preferences}"
    Текст: {{text}}
    Полезный ответ:"""

    map_prompt = PromptTemplate.from_template(map_template)
    summaries = []
    for doc in docs:
        formatted_prompt = map_prompt.format(text=doc.page_content)
        response = client_4.invoke(formatted_prompt)
        summaries.append(response.content)

    combined_summaries = "\n\n".join(summaries)
    reduce_template = f"""Следующий текст состоит из нескольких кратких итогов:
    {{text}}

    На основе этих кратких итогов, выполните финальное сжатие текста, объединяя основные темы и ключевые моменты.
    Уровень сжатия: {compression_percentage}%.
    С учетом пожеланий пользователя "{preferences}"
    Результат предоставьте на русском языке в формате Markdown.

    Полезный ответ:"""

    reduce_prompt = PromptTemplate.from_template(reduce_template)
    final_response = client_4.invoke(reduce_prompt.format(text=combined_summaries))
    return final_response.content


def ask_question_to_mistral_with_large_text(text, question, context, images=[]):
    text_splitter = CharacterTextSplitter.from_huggingface_tokenizer(
        tokenizer, chunk_size=100000, chunk_overlap=14000
    )
    docs = text_splitter.create_documents([text])

    map_template = f"""На основе текста ответьте на вопрос. Ответ должен быть на русском языке.
    Вопрос: {question}
    Текст: {{text}}
    Информация из интернета: {context}
    Полезный ответ:"""

    map_prompt = PromptTemplate.from_template(map_template)
    partial_answers = []
    for doc in docs:
        formatted_prompt = map_prompt.format(text=doc.page_content)
        response = client_4.invoke(formatted_prompt)
        partial_answers.append(response.content)

    combined_answers = "\n\n".join(partial_answers)
    reduce_template = """Следующий текст содержит несколько кратких ответов на вопрос:
    {{text}}

    Объедините их в финальный ответ. Ответ предоставьте на русском языке в формате Markdown.

    Полезный ответ:"""

    reduce_prompt = PromptTemplate.from_template(reduce_template)
    final_response = client_4.invoke(reduce_prompt.format(text=combined_answers))
    return final_response.content


def init_summ(text, preferences, images=[], compression_percentage=30):
    if len(images) >= 8:
        images = images[:7]

    if count_tokens_in_text(text=text) < 128_000:
        return process_article_for_summary(text, preferences, images, compression_percentage)
    else:
        return process_large_article_for_summary(text, preferences, images, compression_percentage)


def init_qa(text, question, images=[]):
    if len(images) >= 8:
        images = images[:7]

    search_tool, tool = setup_search(question)
    context = ''
    if search_tool:
        if tool == 'tavily_tool':
            for result in search_tool:
                context += f"{result.get('url', 'N/A')} : {result.get('content', 'No content')} \n"
        elif tool == 'jina_tool':
            for result in search_tool:
                context += f"{result.get('link', 'N/A')} : {result.get('snippet', 'No snippet')} : {result.get('content', 'No content')} \n"

    if count_tokens_in_text(text=(text + context)) < 128_000:
        return ask_question_to_mistral(text, question, context, images)
    else:
        return ask_question_to_mistral_with_large_text(text, question, context, images)