import time
import json
import logging
from typing import List, Tuple, Generator, Optional, Any
import groq

# Конфигурация
DEFAULT_ITER_MODEL = "llama3-70b-8192"
DEFAULT_FINAL_MODEL = "mixtral-8x7b-32768"
MAX_TOKENS_LIMIT = 8192
RETRY_DELAY = 1.5
MAX_RETRIES = 3

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class APIError(Exception):
    """Базовое исключение для API ошибок"""
    pass

class APIConnectionError(APIError):
    """Ошибка подключения к API"""
    pass

def adjust_max_tokens(model: str, max_tokens: int) -> int:
    """Автоматическая корректировка max_tokens в зависимости от модели"""
    model_limits = {
        "llama3-8b-8192": 8192,
        "llama3-70b-8192": 8192,
        "mixtral-8x7b-32768": 32768
    }
    return min(max_tokens, model_limits.get(model, 4096))

def parse_json_response(raw_response: str) -> dict:
    """Улучшенный парсер с полной валидацией"""
    try:
        data = json.loads(raw_response)
        assert isinstance(data, dict), "Ответ должен быть объектом"
        assert all(k in data for k in ["title", "content", "next_action"]), "Отсутствуют обязательные ключи"
        assert data["next_action"] in {"continue", "final_answer"}, "Некорректное действие"
        return data
    except (json.JSONDecodeError, AssertionError) as e:
        logger.error(f"Ошибка парсинга: {str(e)}")
        return {
            "title": "Ошибка парсинга",
            "content": str(e),
            "next_action": "final_answer"
        }

def make_api_call(
    messages: List[dict],
    max_tokens: int,
    is_final_answer: bool = False,
    custom_client: Optional[Any] = None,
    iter_model: str = DEFAULT_ITER_MODEL,
    final_model: str = DEFAULT_FINAL_MODEL
) -> Any:
    """Улучшенный метод вызова API"""
    client = custom_client or groq.Client()
    model = final_model if is_final_answer else iter_model
    max_tokens = adjust_max_tokens(model, max_tokens)
    
    params = {
        "model": model,
        "messages": messages,
        "max_tokens": max_tokens,
        "temperature": 0.5,
        "response_format": {"type": "json_object"} if not is_final_answer else None
    }

    for attempt in range(1, MAX_RETRIES + 1):
        try:
            response = client.chat.completions.create(**params)
            content = response.choices[0].message.content
            
            if response.choices[0].finish_reason == "length":
                logger.warning(f"Ответ обрезан! Рекомендуемый max_tokens: {max_tokens * 2}")
                
            return content if is_final_answer else parse_json_response(content)
            
        except groq.APIConnectionError as e:
            logger.error(f"Сетевая ошибка (попытка {attempt}/{MAX_RETRIES}): {str(e)}")
            if attempt == MAX_RETRIES:
                raise APIConnectionError(str(e))
            time.sleep(RETRY_DELAY * (2 ** attempt))
            
        except groq.APIError as e:
            logger.error(f"Ошибка API (код {e.status_code}): {e.message}")
            if e.status_code >= 500:
                time.sleep(RETRY_DELAY * attempt)
            else:
                raise APIError(e.message)

def generate_response(
    prompt: str,
    context: Optional[str] = None,
    custom_client: Optional[Any] = None,
    max_steps: int = 10,
    iter_model: str = DEFAULT_ITER_MODEL,
    final_model: str = DEFAULT_FINAL_MODEL
) -> Generator[Tuple[List[Tuple[str, str, float]], Optional[float]], None, None]:
    """Улучшенная генерация ответа"""
    system_prompt = (
        "Вы — AI-ассистент для анализа технических вопросов. "
        "Формат ответа: JSON с ключами title, content, next_action. "
        "Используйте русский язык и технические термины."
    )
    
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": f"Контекст: {context}\n\nВопрос: {prompt}" if context else prompt}
    ]
    
    steps = []
    total_time = 0.0
    
    try:
        for step in range(1, max_steps + 1):
            start = time.monotonic()
            
            try:
                response = make_api_call(
                    messages=messages,
                    max_tokens=MAX_TOKENS_LIMIT,
                    is_final_answer=False,
                    custom_client=custom_client,
                    iter_model=iter_model
                )
            except APIError as e:
                logger.error("Критическая ошибка: %s", str(e))
                yield [("Ошибка", str(e), time.monotonic() - start)], None
                return
                
            elapsed = time.monotonic() - start
            total_time += elapsed
            
            steps.append((f"Шаг {step}: {response['title']}", response["content"], elapsed))
            messages.append({"role": "assistant", "content": json.dumps(response, ensure_ascii=False)})
            
            if response["next_action"] == "final_answer":
                break
                
            yield steps, None
            
        # Финализация
        messages.append({"role": "user", "content": "Сформируйте финальный ответ с примерами кода и пояснениями."})
        
        final_start = time.monotonic()
        final_answer = make_api_call(
            messages=messages,
            max_tokens=MAX_TOKENS_LIMIT,
            is_final_answer=True,
            custom_client=custom_client,
            final_model=final_model
        )
        total_time += time.monotonic() - final_start
        
        steps.append(("Финальный ответ", final_answer, time.monotonic() - final_start))
        
    except Exception as e:
        logger.exception("Непредвиденная ошибка:")
        steps.append(("Критическая ошибка", str(e), 0.0))
        
    yield steps, total_time

if __name__ == "__main__":
    try:
        for steps, total_time in generate_response(
            prompt="Опишите процесс обработки запросов в Groq API",
            context="Используйте официальную документацию Groq",
            iter_model="llama3-70b-8192",
            final_model="mixtral-8x7b-32768"
        ):
            if total_time is not None:
                print(f"\n✅ Готово за {total_time:.2f}s")
                print("=" * 60)
            for title, content, t in steps:
                print(f"\n🔹 {title} [{t:.2f}s]\n{content}")
                
    except APIConnectionError as e:
        print(f"🚨 Ошибка подключения: {str(e)}")
    except APIError as e:
        print(f"🚨 Ошибка API: {str(e)}")