import groq
import time
import os
import json
import logging

logging.basicConfig(level=logging.DEBUG)

client = groq.Groq()

def parse_json_response(raw_response):
    """
    Разбирает строку JSON и проверяет наличие обязательных ключей: title, content, next_action.
    """
    logging.debug(f"RAW JSON RESPONSE: {raw_response}")
    try:
        parsed_response = json.loads(raw_response)
    except json.JSONDecodeError as e:
        raise ValueError(f"Invalid JSON: {e}, content: {raw_response}")
    
    required_keys = {"title", "content", "next_action"}
    if required_keys.issubset(parsed_response):
        return parsed_response
    else:
        raise ValueError("JSON missing required keys")

def make_api_call(messages, max_tokens, is_final_answer=False, custom_client=None):
    """
    Отправляет API-запрос с заданными сообщениями и возвращает результат.
    
    При is_final_answer=True возвращает текст финального ответа.
    В противном случае ожидается, что ответ будет в формате JSON с ключами "title", "content", "next_action".
    """
    current_client = custom_client if custom_client is not None else client

    for attempt in range(3):
        try:
            if is_final_answer:
                response = current_client.chat.completions.create(
                    model="llama3-8b-8192",
                    messages=messages,
                    max_tokens=max_tokens,
                    temperature=0.2
                )
                # Возвращаем строку, как и ожидается для финального ответа
                return response.choices[0].message.content
            else:
                response = current_client.chat.completions.create(
                    model="llama3-8b-8192",
                    messages=messages,
                    max_tokens=max_tokens,
                    temperature=0.2,
                    response_format={"type": "json_object"}
                )
                raw_response = response.choices[0].message.content
                return parse_json_response(raw_response)
        except Exception as e:
            logging.exception("Attempt %d failed", attempt + 1)
            if attempt == 2:
                error_message = f"Failed to generate {'final answer' if is_final_answer else 'step'} after 3 attempts. Error: {str(e)}"
                if is_final_answer:
                    return f"Error: {error_message}"
                else:
                    return {
                        "title": "Error",
                        "content": error_message,
                        "next_action": "final_answer"
                    }
            time.sleep(1)

def generate_response(prompt, custom_client=None):
    """
    Генерирует цепочку рассуждений и финальный ответ по заданному запросу.
    
    Аргументы:
        prompt (str): Запрос, на который необходимо сгенерировать ответ.
        custom_client (object, optional): Альтернативный клиент для API-вызовов.
    
    Возвращает:
        Генератор, yield'ящий кортеж (steps, total_thinking_time), где:
            - steps: список кортежей (название шага, содержание, время обработки)
            - total_thinking_time: общее время обработки (либо None до финального шага)
    """
    messages = [
        {"role": "system", "content": (
            "Вы – интеллектуальный помощник, который анализирует и объясняет свои рассуждения на русском языке шаг за шагом.\n"
            "### 🔹 Формат ответа\n"
            "Ваш ответ должен быть строго в JSON-формате без дополнительного текста или форматирования (например, без ```json```).\n"
            "Обязательные ключи:\n"
            '- "title" – краткое название шага.\n'
            '- "content" – описание действий.\n'
            '- "next_action" – "continue" или "final_answer".\n'
            "Пример:\n"
            '{"title": "Анализ задачи", "content": "Выделение ключевых элементов...", "next_action": "continue"}\n'
            "🔹 Дополнительные требования:\n"
            "- Используйте русский язык.\n"
            '- Избегайте Unicode-кодировок (например, писать "Привет", а не "\\u041f\\u0440\\u0438...").'
        )}, 
        {"role": "user", "content": prompt},
        {"role": "assistant", "content": "Спасибо! Начинаю анализ..."}
    ]

    steps = []
    step_count = 1
    total_thinking_time = 0

    while True:
        start_time = time.time()
        step_data = make_api_call(messages, max_tokens=500, custom_client=custom_client)
        end_time = time.time()
        thinking_time = end_time - start_time
        total_thinking_time += thinking_time

        steps.append((f"Step {step_count}: {step_data.get('title', 'Без названия')}", step_data.get('content', ''), thinking_time))
        messages.append({"role": "assistant", "content": json.dumps(step_data, ensure_ascii=False)})

        if step_data.get('next_action') == 'final_answer' or step_count >= 25:
            break

        step_count += 1
        yield steps, None  # Возвращаем промежуточные шаги без финального времени

    messages.append({
        "role": "user",
        "content": "Предоставьте окончательный ответ без формата JSON. Сохранить исходное форматирование из подсказки."
    })

    start_time = time.time()
    final_data = make_api_call(messages, max_tokens=1200, is_final_answer=True, custom_client=custom_client)
    end_time = time.time()
    thinking_time = end_time - start_time
    total_thinking_time += thinking_time

    steps.append(("Final Answer", final_data, thinking_time))
    yield steps, total_thinking_time