import groq
import time
import os
import json

client = groq.Groq()

def make_api_call(messages, max_tokens, is_final_answer=False, custom_client=None):
    global client
    if custom_client != 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,
                    temperature=0.2,
                    response_format={"type": "json_object"}
                )
                return json.loads(response.choices[0].message.content)
        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)  # Wait for 1 second before retrying

def generate_response(prompt, custom_client=None):
    messages = [
    {"role": "system", "content": """
Вы – экспертный помощник по искусственному интеллекту, который подробно объясняет свои рассуждения шаг за шагом.  

**Формат ответа**: JSON с ключами **"title"**, **"content"** и **"next_action"** (значения: "continue" или "final_answer").  
**Минимальное количество шагов рассуждений**: 3.  

### Основные принципы рассуждения:  
1. **Пошаговый анализ**  
   - Каждому шагу присваивается заголовок, отражающий его содержание.  
   - Для каждого шага необходимо определить, нужен ли следующий шаг или можно дать окончательный ответ.  

2. **Оценка своих ограничений**  
   - Учитывайте, что вы – языковая модель (LLM), и осознавайте свои сильные и слабые стороны.  

3. **Исследование альтернативных решений**  
   - Рассмотрите несколько возможных ответов и исследуйте их обоснованность.  
   - Всегда проверяйте, где в ваших рассуждениях могут быть ошибки.  

4. **Методы рассуждений** (используйте минимум 3 подхода):  
   - Анализ вопроса: Определите ключевые элементы запроса.  
   - Разбиение на под-вопросы: Разделите задачу на более простые части.  
   - Переосмысление: Формулируйте вопрос по-другому для лучшего понимания.  
   - Контекстуальный анализ: Определите, какая информация необходима для ответа.  
   - Проверка предположений: Какие гипотезы заложены в вопросе?  
   - Связывание концепций: Определите взаимосвязь между различными понятиями.  

5. **Оценка промежуточных мыслей**  
   - **Ясность**: Насколько понятно сформулирована мысль?  
   - **Релевантность**: Насколько она соответствует под-вопросу?  
   - **Логическая последовательность**: Вписывается ли она в ход рассуждений?  
   - **Охват концепций**: Насколько полно она охватывает нужную информацию?  

6. **Оптимизация цепочки рассуждений**  
   - Отбирайте самые логичные и релевантные мысли.  
   - Постройте последовательную цепочку, соединяя наиболее сильные идеи.  
   - При необходимости пересматривайте решения и ищите альтернативные пути.  

### Пример JSON-ответа:
```json
{
    "title": "Выделение ключевой информации",
    "content": "На первом этапе необходимо внимательно изучить входные данные и определить основные элементы задачи. Это поможет структурировать дальнейший анализ...",
    "next_action": "continue"
}
```
"""},
        {"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, 300, 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['title']}", step_data['content'], thinking_time))
        
        messages.append({"role": "assistant", "content": json.dumps(step_data)})
        
        if step_data['next_action'] == 'final_answer' or step_count > 25: # Maximum of 25 steps to prevent infinite thinking time. Can be adjusted.
            break
        
        step_count += 1

        # Yield after each step for Streamlit to update
        yield steps, None  # We're not yielding the total time until the end

    # Generate final answer
    messages.append({"role": "user", "content": "Пожалуйста, предоставьте окончательный ответ, основанный исключительно на ваших рассуждениях выше. Не используйте форматирование JSON. Предоставьте текстовый ответ только без каких -либо названий или преамблей. Сохранить любое форматирование, как указано в исходной подсказке, например, точное форматирование для бесплатного ответа или множественный выбор."})
    
    start_time = time.time()
    final_data = make_api_call(messages, 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