from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage, SystemMessage
from state import ReflectState

# Инициализация LLM
llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0.7)

# Промпты для узлов
DRAFT_PROMPT = """Ты — эксперт, который дает краткие, но содержательные ответы.
Ответь на вопрос пользователя в 5-10 предложений.
Будь конкретным, избегай общих фраз и "воды".

Вопрос: {question}"""

REFLECT_PROMPT = """Ты — строгий критик, оценивающий качество ответа.
Критерии оценки:
1. Полнота: ответ покрывает все аспекты вопроса
2. Конкретика: есть примеры, детали, точные формулировки
3. Отсутствие воды: нет общих фраз, повторений, лишних слов

Ответ для оценки: {draft}

Твой ответ должен быть в формате:
VERDICT: [ok/needs_revision]
CRITIQUE: [2-3 конкретных пункта замечаний или похвала]

Примеры:
VERDICT: ok
CRITIQUE: Ответ полный, есть примеры, нет воды.

VERDICT: needs_revision
CRITIQUE: 1. Нет конкретных примеров. 2. Слишком много общих фраз. 3. Не раскрыта разница в контексте MCP."""

REWRITE_PROMPT = """Улучши ответ на вопрос, учитывая замечания критика.
Вопрос: {question}
Предыдущая версия: {draft}
Замечания критика: {critique}

Дай исправленную версию (5-10 предложений), которая:
- Учитывает все замечания
- Становится более полной и конкретной
- Убирает воду
- Сохраняет краткость (5-10 предложений)"""

def draft_answer(state: ReflectState) -> ReflectState:
    """
    Узел: первичная генерация ответа с try/except
    При ошибке заполняет поле error
    """
    print(f"\n📝 [Раунд {state['round']}] Генерация черновика...")
    
    try:
        prompt = DRAFT_PROMPT.format(question=state["question"])
        response = llm.invoke([HumanMessage(content=prompt)])
        
        # Проверка качества ответа (базовая валидация)
        if len(response.content.split()) < 50:
            raise ValueError("Слишком короткий ответ (<50 слов)")
        
        if len(response.content.split()) > 300:
            raise ValueError("Слишком длинный ответ (>300 слов)")
        
        print(f"   ✅ Черновик успешно создан")
        new_state = state.copy()
        new_state["draft"] = response.content
        new_state["error"] = None  # Очищаем ошибку
        return new_state
        
    except Exception as e:
        print(f"   ❌ Ошибка генерации: {str(e)}")
        new_state = state.copy()
        new_state["draft"] = ""
        new_state["error"] = str(e)
        return new_state

def reflect(state: ReflectState) -> ReflectState:
    """Узел: критическая оценка ответа LLM-критиком"""
    print(f"🔍 [Раунд {state['round']}] Рефлексия критика...")
    
    try:
        prompt = REFLECT_PROMPT.format(draft=state["draft"])
        response = llm.invoke([HumanMessage(content=prompt)])
        
        # Парсим ответ критика
        lines = response.content.strip().split('\n')
        verdict = "needs_revision"
        critique = response.content
        
        for line in lines:
            if line.startswith("VERDICT:"):
                verdict = line.replace("VERDICT:", "").strip().lower()
            elif line.startswith("CRITIQUE:"):
                critique = line.replace("CRITIQUE:", "").strip()
        
        if verdict not in ["ok", "needs_revision"]:
            verdict = "needs_revision"
        
        new_state = state.copy()
        new_state["verdict"] = verdict
        new_state["critique"] = critique
        
        print(f"   Вердикт: {verdict.upper()}")
        if verdict == "needs_revision":
            print(f"   Замечания: {critique[:100]}...")
        
        return new_state
        
    except Exception as e:
        # Если критик сломался, считаем что нужна доработка
        print(f"   ⚠️ Ошибка критика: {str(e)}")
        new_state = state.copy()
        new_state["verdict"] = "needs_revision"
        new_state["critique"] = f"Критик не смог оценить ответ: {str(e)}"
        return new_state

def rewrite(state: ReflectState) -> ReflectState:
    """Узел: переписывание ответа с учетом замечаний"""
    print(f"✏️ [Раунд {state['round']}] Переписывание с учетом критики...")
    
    try:
        prompt = REWRITE_PROMPT.format(
            question=state["question"],
            draft=state["draft"],
            critique=state["critique"]
        )
        response = llm.invoke([HumanMessage(content=prompt)])
        
        new_state = state.copy()
        new_state["draft"] = response.content
        new_state["round"] = state["round"] + 1
        new_state["error"] = None
        return new_state
        
    except Exception as e:
        print(f"   ❌ Ошибка переписывания: {str(e)}")
        new_state = state.copy()
        new_state["error"] = str(e)
        new_state["round"] = state["round"] + 1
        return new_state