# search_article.py (адаптировано под LangChain >= 1.0)
import os
import json
import concurrent.futures
from mistralai.client import Mistral
import arxiv
from transformers import AutoTokenizer
from langchain_core.prompts import PromptTemplate
from langchain_text_splitters import CharacterTextSplitter
from langchain_mistralai import ChatMistralAI

api_key_1 = os.getenv("MISTRAL_API_KEY_1")
client_1 = Mistral(api_key=api_key_1)

llm = ChatMistralAI(api_key=api_key_1, model="pixtral-12b-2409")

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

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 extract_key_topics(content, images=[]):
    prompt = f"""
    Extract the primary themes from the text below. List each theme in as few words as possible, focusing on essential concepts only. Format as a concise, unordered list with no extraneous words.

    ```{content}```

    LIST IN ENGLISH:
    -
    """
    message_content = [{"type": "text", "text": prompt}] + images

    response = client_1.chat.complete(
        model="pixtral-12b-2409",
        messages=[{"role": "user", "content": message_content}]
    )
    return response.choices[0].message.content


def extract_key_topics_with_large_text(content, images=[]):
    text_splitter = CharacterTextSplitter.from_huggingface_tokenizer(
        tokenizer, chunk_size=100000, chunk_overlap=14000
    )
    docs = text_splitter.create_documents([content])

    image_descriptions = "\n".join(
        [f"Изображение {i+1}: {img['image_url']}" for i, img in enumerate(images)]
    )

    map_template = f"""
        Текст: {{text}}
        Изображения: {image_descriptions}

        Extract the primary themes from the text below. List each theme in as few words as possible, focusing on essential concepts only. Format as a concise, unordered list with no extraneous words.
        LIST IN ENGLISH:
        -
    """

    map_prompt = PromptTemplate.from_template(map_template)
    chunk_themes = []

    for doc in docs:
        formatted_prompt = map_prompt.format(text=doc.page_content)
        response = llm.invoke(formatted_prompt)
        chunk_themes.append(response.content)

    combined_themes = "\n".join(chunk_themes)

    reduce_template = f"""
        Следующий текст содержит несколько списков ключевых тем, извлеченных из разных частей документа:
        {{themes}}

        Extract the primary themes from the text below. List each theme in as few words as possible, focusing on essential concepts only. Format as a concise, unordered list with no extraneous words.
        Remove duplicates and merge similar themes.

        LIST IN ENGLISH:
        -
    """

    reduce_prompt = PromptTemplate.from_template(reduce_template)
    final_response = llm.invoke(reduce_prompt.format(themes=combined_themes))

    return final_response.content


def search_relevant_articles_arxiv(key_topics, max_articles=10):
    articles_by_topic = {}
    final_topics = []

    def fetch_articles_for_topic(topic):
        topic_articles = []
        try:
            search = arxiv.Search(
                query=topic,
                max_results=max_articles,
                sort_by=arxiv.SortCriterion.Relevance
            )
            for result in search.results():
                article_data = {
                    "title": result.title,
                    "doi": result.doi,
                    "summary": result.summary,
                    "url": result.entry_id,
                    "pdf_url": result.pdf_url
                }
                topic_articles.append(article_data)
            final_topics.append(topic)
        except Exception as e:
            print(f"Error fetching articles for topic '{topic}': {e}")

        return topic, topic_articles

    with concurrent.futures.ThreadPoolExecutor(max_workers=10) as executor:
        futures = {executor.submit(fetch_articles_for_topic, topic): topic for topic in key_topics}
        for future in concurrent.futures.as_completed(futures):
            topic, articles = future.result()
            if articles:
                articles_by_topic[topic] = articles

    return articles_by_topic, list(set(final_topics))


def init(content, images=[]):
    if len(images) >= 8:
        images = images[:7]

    if count_tokens_in_text(content) < 128000:
        key_topics = extract_key_topics(content, images)
    else:
        key_topics = extract_key_topics_with_large_text(content, images)

    key_topics = [topic.strip("- ") for topic in key_topics.split("\n") if topic.strip()]

    articles_by_topic, final_topics = search_relevant_articles_arxiv(key_topics)

    result_json = json.dumps(articles_by_topic, indent=4)

    return final_topics, result_json