# utils/agent_runner.py
import json

def run_agent_safely(agent, max_retries=3, return_string=False, **kwargs):
    """
    Безопасно запускает агента с повторными попытками при ошибках.
    
    Args:
        agent: Экземпляр агента smolagents
        max_retries: Максимальное количество попыток
        return_string: Если True, возвращает строку вместо парсинга JSON
        **kwargs: Аргументы для передачи агенту
        
    Returns:
        dict или str: Распарсенный JSON ответ или строка (в зависимости от return_string)
    """
    last_error = None

    for attempt in range(1, max_retries + 1):
        try:
            result = agent.run(**kwargs)

            # Извлекаем ответ из результата smolagents
            # smolagents возвращает AgentText объект, который нужно правильно обработать
            raw = None
            
            # Если это AgentText объект (из smolagents), получаем полный текст
            # AgentText может содержать обрезанный вывод при str(), но полный текст доступен через другие методы
            if hasattr(result, 'to_string'):
                raw = result.to_string()
            elif hasattr(result, 'text'):
                raw = result.text
            elif hasattr(result, '__getitem__'):
                # Если это последовательность или маппинг, пробуем извлечь
                try:
                    if 'answer' in result:
                        raw = result['answer']
                except:
                    pass
            elif isinstance(result, dict):
                # Ищем ответ в различных возможных полях
                if "answer" in result:
                    raw = result["answer"]
                elif "content" in result:
                    raw = result["content"]
                elif "text" in result:
                    raw = result["text"]
                elif "response" in result:
                    raw = result["response"]
                else:
                    # Если словарь не содержит текстовых полей, преобразуем в JSON строку
                    raw = json.dumps(result, ensure_ascii=False)
            elif isinstance(result, str):
                raw = result
            else:
                # Пробуем преобразовать в строку
                raw = str(result)

            if not isinstance(raw, str):
                raise ValueError(f"Agent output is not a string, got {type(raw)}")
            
            # Удаляем возможные escape-последовательности и пробелы
            raw = raw.strip()
            
            # Если raw содержит JSON-строку внутри (например, {"answer": "{\"key\": \"value\"}"})
            # Нужно распарсить внешний JSON, чтобы получить внутренний
            if raw.startswith('{"') and '"answer"' in raw:
                try:
                    # Пробуем распарсить как JSON
                    outer_json = json.loads(raw)
                    if isinstance(outer_json, dict) and "answer" in outer_json:
                        raw = outer_json["answer"]
                except json.JSONDecodeError:
                    # Если не получается распарсить, возможно JSON обрезан
                    # Пробуем найти полный JSON внутри
                    pass
            
            # Если нужно вернуть строку (для GPT Prompt агента)
            if return_string:
                return raw
            
            # Для JSON агентов - пытаемся распарсить
            # Если строка обрезана (начинается с { но не заканчивается }), пробуем исправить
            if raw.startswith('{') and not raw.rstrip().endswith('}'):
                # Пытаемся найти последнюю закрывающую скобку
                last_brace = raw.rfind('}')
                if last_brace > len(raw) * 0.8:  # Если скобка не слишком близко к началу
                    raw = raw[:last_brace + 1]
            
            # Парсим JSON
            try:
                return json.loads(raw)
            except json.JSONDecodeError as e:
                # Если это не JSON, пробуем найти JSON внутри строки
                # Ищем первую { и последнюю }
                start_idx = raw.find('{')
                end_idx = raw.rfind('}')
                if start_idx >= 0 and end_idx > start_idx:
                    json_part = raw[start_idx:end_idx + 1]
                    try:
                        return json.loads(json_part)
                    except:
                        pass
                
                # Если всё равно не получается, это ошибка для JSON-агентов
                if not raw.startswith('{') and not raw.startswith('['):
                    raise ValueError(f"Expected JSON but got plain text (first 200 chars): {raw[:200]}")
                raise ValueError(f"JSON decode error: {e}. Raw (first 500 chars): {raw[:500]}")

        except Exception as e:
            last_error = e
            print(f"[{agent.name}] Attempt {attempt}/{max_retries} failed: {e}")

    raise ValueError(
        f"[{agent.name}] Failed after {max_retries} attempts.\nLast error: {last_error}"
    )