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.3)

# Инициализация Tavily search tool
tavily_search = TavilySearch(
    max_results=3,
    topic="general",
    include_answer=False,
    include_raw_content=False
)

# Промпты для узлов
OUTLINE_PROMPT = """Ты — исследовательский ассистент. Составь план из 4-5 пунктов для исследования темы.

Тема: {topic}

Требования к плану:
- Каждый пункт должен быть конкретным и исследовательским
- Пункты должны покрывать разные аспекты темы
- Используй четкие формулировки (например: "История развития", "Ключевые технологии", "Практические примеры")

Ответ дай в виде списка, каждый пункт на новой строке, начиная с "- ".

Пример:
- Текущее состояние и тренды
- Основные вызовы и решения
- Рекомендации для начинающих"""

RESEARCH_PROMPT = """Ты — исследователь. На основе результатов поиска составь краткую заметку (5-8 предложений) по следующему пункту плана.

Пункт плана: {outline_item}

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

Заметка должна:
- Быть информативной и конкретной
- Опираться на факты из поиска
- Содержать 5-8 предложений
- Заканчиваться выводом или ключевым инсайтом"""

SYNTHESIS_PROMPT = """Ты — аналитик. Объедини собранные заметки в связный исследовательский бриф (полстраницы - страница).

Тема: {topic}

Заметки по пунктам плана:
{notes}

Требования к брифу:
1. Начни с краткого введения (2-3 предложения)
2. Каждый пункт плана оформи как отдельный раздел с заголовком
3. Сохрани ключевые факты и инсайты из заметок
4. Добавь заключение с общими выводами (3-4 предложения)
5. Используй связующие фразы между разделами

Формат:
# Исследовательский бриф: {topic}

## Введение
[текст]

## [Заголовок пункта 1]
[текст]

## [Заголовок пункта 2]
[текст]

...

## Заключение
[текст]"""

def outline_node(state: BriefState) -> BriefState:
    """Узел 1: Составление плана исследования"""
    print("\n📋 [1/4] Составляю план исследования...")
    
    prompt = OUTLINE_PROMPT.format(topic=state["topic"])
    response = llm.invoke([HumanMessage(content=prompt)])
    
    # Парсим список из ответа
    lines = response.content.strip().split('\n')
    outline = []
    for line in lines:
        line = line.strip()
        if line.startswith('- '):
            outline.append(line[2:])
        elif line.startswith('• '):
            outline.append(line[2:])
        elif line and not line.startswith('```'):
            # Если нет маркера, убираем номер
            clean = line.lstrip('0123456789. ').strip()
            if clean:
                outline.append(clean)
    
    # Ограничиваем 5 пунктами
    outline = outline[:5]
    
    print(f"   ✅ План составлен: {len(outline)} пунктов")
    for i, item in enumerate(outline, 1):
        print(f"      {i}. {item}")
    
    new_state = state.copy()
    new_state["outline"] = outline
    new_state["step_index"] = 0
    new_state["notes"] = []
    return new_state

def research_step_node(state: BriefState) -> BriefState:
    """Узел 2: Исследование одного пункта плана через web search"""
    current_idx = state["step_index"]
    current_item = state["outline"][current_idx]
    
    print(f"\n🔍 [2/4] Исследую пункт {current_idx + 1}/{len(state['outline'])}: {current_item}")
    
    # Формируем поисковый запрос
    search_query = f"{state['topic']} {current_item}"
    print(f"   Запрос в Tavily: {search_query}")
    
    try:
        # Выполняем поиск
        search_result = tavily_search.invoke({"query": search_query})
        
        # Извлекаем текст результатов
        if isinstance(search_result, list):
            results_text = "\n\n".join([
                f"Источник {i+1}: {item.get('content', '')[:500]}"
                for i, item in enumerate(search_result[:3])
                if item.get('content')
            ])
        else:
            results_text = str(search_result)[:1500]
        
        # Генерируем заметку на основе поиска
        prompt = RESEARCH_PROMPT.format(
            outline_item=current_item,
            search_results=results_text
        )
        response = llm.invoke([HumanMessage(content=prompt)])
        note = response.content.strip()
        
        print(f"   ✅ Заметка создана ({len(note)} символов)")
        
    except Exception as e:
        print(f"   ⚠️ Ошибка поиска: {e}")
        # Fallback: заметка без поиска
        prompt = f"Напиши краткую заметку (5-8 предложений) о: {current_item} в контексте {state['topic']}"
        response = llm.invoke([HumanMessage(content=prompt)])
        note = response.content.strip()
        print(f"   ✅ Заметка создана (без поиска)")
    
    # Обновляем состояние
    new_state = state.copy()
    new_state["notes"] = state["notes"] + [note]
    new_state["step_index"] = state["step_index"] + 1
    
    return new_state

def synthesize_node(state: BriefState) -> BriefState:
    """Узел 3: Синтез финального брифа из заметок"""
    print("\n📝 [3/4] Синтезирую финальный исследовательский бриф...")
    
    # Форматируем заметки с заголовками пунктов
    notes_with_headers = []
    for i, (outline_item, note) in enumerate(zip(state["outline"], state["notes"])):
        notes_with_headers.append(f"## {outline_item}\n{note}")
    
    notes_text = "\n\n".join(notes_with_headers)
    
    prompt = SYNTHESIS_PROMPT.format(
        topic=state["topic"],
        notes=notes_text
    )
    
    response = llm.invoke([HumanMessage(content=prompt)])
    final_brief = response.content.strip()
    
    print(f"   ✅ Бриф создан ({len(final_brief)} символов)")
    
    new_state = state.copy()
    new_state["final_brief"] = final_brief
    return new_state