import os
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage, SystemMessage
from langchain_tavily import TavilySearch
from state import BriefState

# Инициализация LLM и поискового инструмента
llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0.5)
tavily = TavilySearch(
    api_key=os.getenv("TAVILY_API_KEY"),
    max_results=3,  # Количество результатов поиска
    include_answer=True,
    include_raw_content=False
)

# Промпты
OUTLINE_PROMPT = """Ты — исследовательский ассистент. Составь план исследования по теме: {topic}

Требования:
- 4-5 пунктов
- Каждый пункт должен быть конкретным аспектом для исследования
- Пункты должны логически раскрывать тему
- Формат вывода: просто список пунктов (каждый с новой строки, начинается с цифры)

Пример:
1. Основные понятия и определения
2. Ключевые технологии и инструменты
3. Практические примеры использования
4. Потенциальные риски и их mitigation
5. Лучшие практики и рекомендации

Тема: {topic}
План:"""

RESEARCH_PROMPT = """Ты — исследователь, собирающий информацию по конкретному пункту плана.

Тема исследования: {topic}
Текущий пункт плана: {outline_item}

На основе результатов поиска (ниже) составь краткую заметку (5-8 предложений), которая:
1. Отвечает на пункт плана
2. Содержит конкретные факты, данные, примеры
3. Использует информацию из результатов поиска
4. Пишется связным текстом (не маркированным списком)

Результаты поиска:
{search_results}

Заметка по пункту "{outline_item}":"""

SYNTHESIS_PROMPT = """Ты — аналитик, который создает итоговый исследовательский бриф.

Тема: {topic}
План исследования:
{outline}

Собранные заметки по пунктам:
{notes}

Задание: создай связный исследовательский бриф (объемом ½ - 1 страницу), который:
1. Имеет заголовок "Исследовательский бриф: {topic}"
2. Структурирован по пунктам плана (каждый пункт как подзаголовок)
3. Содержит введение (1 абзац) и заключение (1 абзац)
4. Объединяет заметки в единый поток, убирает повторы
5. Добавляет связующие предложения между пунктами
6. Завершается списком ключевых выводов (3-4 пункта)

Форматирование: используй markdown для заголовков и списков.

Бриф:"""

def outline_node(state: BriefState) -> BriefState:
    """Узел: генерация плана исследования"""
    print(f"\n📋 Генерация плана исследования по теме: {state['topic']}")
    
    prompt = OUTLINE_PROMPT.format(topic=state["topic"])
    response = llm.invoke([HumanMessage(content=prompt)])
    
    # Парсим план в список
    outline_lines = response.content.strip().split('\n')
    outline = []
    for line in outline_lines:
        # Убираем номера пунктов и лишние символы
        cleaned = line.strip()
        if cleaned and cleaned[0].isdigit():
            # Убираем "1. " или "1)" и т.д.
            import re
            cleaned = re.sub(r'^\d+[\.\)]\s*', '', cleaned)
            outline.append(cleaned)
    
    print(f"✅ Создан план из {len(outline)} пунктов:")
    for i, item in enumerate(outline, 1):
        print(f"   {i}. {item}")
    
    return {
        **state,
        "outline": outline,
        "step_index": 0,
        "notes": []
    }

def research_step_node(state: BriefState) -> BriefState:
    """Узел: исследование одного пункта плана с поиском в интернете"""
    current_step = state["step_index"]
    outline_item = state["outline"][current_step]
    
    print(f"\n🔍 [Шаг {current_step + 1}/{len(state['outline'])}] Исследование: {outline_item}")
    
    # Формируем поисковый запрос
    search_query = f"{state['topic']} {outline_item}"
    print(f"   Поисковый запрос: {search_query}")
    
    try:
        # Выполняем поиск через Tavily
        search_result = tavily.invoke({"query": search_query})
        
        # Форматируем результаты поиска
        if isinstance(search_result, list):
            formatted_results = []
            for idx, result in enumerate(search_result, 1):
                if hasattr(result, 'content') and result.content:
                    formatted_results.append(f"{idx}. {result.content[:500]}")
                elif isinstance(result, dict) and 'content' in result:
                    formatted_results.append(f"{idx}. {result['content'][:500]}")
            search_text = "\n\n".join(formatted_results) if formatted_results else "Результаты не найдены."
        elif isinstance(search_result, str):
            search_text = search_result
        else:
            search_text = str(search_result)
        
        # Генерируем заметку на основе поиска
        research_prompt = RESEARCH_PROMPT.format(
            topic=state["topic"],
            outline_item=outline_item,
            search_results=search_text[:1500]  # Ограничиваем длину
        )
        
        response = llm.invoke([HumanMessage(content=research_prompt)])
        note = response.content.strip()
        
        print(f"   ✅ Заметка создана ({len(note)} символов)")
        
    except Exception as e:
        print(f"   ⚠️ Ошибка поиска: {e}")
        # Fallback: генерируем заметку без поиска
        fallback_prompt = f"Напиши краткую заметку (5-8 предложений) по пункту '{outline_item}' для темы '{state['topic']}' (без внешнего поиска)."
        response = llm.invoke([HumanMessage(content=fallback_prompt)])
        note = response.content.strip()
        print(f"   ⚠️ Создана fallback-заметка без поиска")
    
    # Добавляем заметку в список
    updated_notes = state["notes"] + [note]
    
    return {
        **state,
        "notes": updated_notes,
        "step_index": state["step_index"] + 1
    }

def synthesize_node(state: BriefState) -> BriefState:
    """Узел: синтез итогового брифа из собранных заметок"""
    print(f"\n📝 Синтез итогового исследовательского брифа...")
    
    # Форматируем план и заметки для промпта
    outline_text = "\n".join([f"{i+1}. {item}" for i, item in enumerate(state["outline"])])
    notes_text = "\n\n---\n\n".join([
        f"**Пункт {i+1}: {state['outline'][i]}**\n{note}"
        for i, note in enumerate(state["notes"])
    ])
    
    prompt = SYNTHESIS_PROMPT.format(
        topic=state["topic"],
        outline=outline_text,
        notes=notes_text
    )
    
    response = llm.invoke([HumanMessage(content=prompt)])
    final_brief = response.content.strip()
    
    print(f"✅ Бриф создан ({len(final_brief)} символов)")
    
    return {
        **state,
        "final_brief": final_brief
    }