import groq
import time
import os
import json
import logging

logging.basicConfig(level=logging.DEBUG)

client = groq.Groq()

def make_api_call(messages, max_tokens, is_final_answer=False, custom_client=None):
    global client
    if custom_client is not None:
        client = custom_client

    for attempt in range(3):
        try:
            if is_final_answer:
                response = 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 = client.chat.completions.create(
                    model="llama3-8b-8192",
                    messages=messages,
                    max_tokens=max_tokens,  # Используем переданный параметр max_tokens
                    temperature=0.2,
                    response_format={"type": "json_object"}
                )

                raw_response = response.choices[0].message.content
                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}")

                if all(k in parsed_response for k in ("title", "content", "next_action")):
                    return parsed_response
                else:
                    raise ValueError("JSON missing required keys")
        except Exception as e:
            if attempt == 2:
                if is_final_answer:
                    return {
                        "title": "Error",
                        "content": f"Failed to generate final answer after 3 attempts. Error: {str(e)}"
                    }
                else:
                    return {
                        "title": "Error",
                        "content": f"Failed to generate step after 3 attempts. Error: {str(e)}",
                        "next_action": "final_answer"
                    }
            time.sleep(1)

def generate_response(prompt, custom_client=None):
    """
    Генерирует ответ, возвращая промежуточные шаги рассуждений и финальный результат.
    Функция является генератором: yield возвращает кортеж (steps, total_thinking_time).
    """
    messages = [
        {"role": "system", "content": """
Вы – интеллектуальный помощник, который анализирует и объясняет свои рассуждения на русском языке шаг за шагом.
### 🔹 Формат ответа
Ваш ответ должен быть строго в JSON-формате без дополнительного текста или форматирования (например, без ```json```).
Обязательные ключи:
- "title" – краткое название шага.
- "content" – описание действий.
- "next_action" – "continue" или "final_answer".
Пример:
{"title": "Анализ задачи", "content": "Выделение ключевых элементов...", "next_action": "continue"}
🔹 Дополнительные требования:
- Используйте русский язык.
- Избегайте 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)  # Передаём max_tokens
        end_time = time.time()
        thinking_time = end_time - start_time
        total_thinking_time += thinking_time

        steps.append((f"Step {step_count}: {step_data['title']}", step_data['content'], thinking_time))
        messages.append({"role": "assistant", "content": json.dumps(step_data)})

        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