# agents/coordinator.py
import logging
import json
import re
from typing import Dict, Any, List
from agents.nebius_simple import create_nebius_llm

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


class CoordinatorAgent:
    def __init__(self, nebius_api_key: str):
        self.llm = create_nebius_llm(
            api_key=nebius_api_key,
            model="meta-llama/Llama-3.3-70B-Instruct-fast",
            temperature=0.6  # Чуть выше для креативности при генерации историй
        )

    def _extract_json(self, text: str) -> dict:
        """
        Извлекает JSON из ответа LLM с улучшенной обработкой.
        Поддерживает различные форматы ответов.
        """
        # 1. Убираем markdown code blocks
        text = re.sub(r"```json\s*", "", text)
        text = re.sub(r"```\s*", "", text)

        # 2. Ищем JSON объект с помощью regex
        json_pattern = r'\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}'
        matches = re.findall(json_pattern, text, re.DOTALL)

        if matches:
        # Берём первый найденный JSON объект
            json_str = matches[0].strip()
        try:
            return json.loads(json_str)
        except json.JSONDecodeError as e:
            logger.warning(f"Failed to parse extracted JSON: {e}")
            logger.debug(f"JSON string was: {json_str}")

        # 3. Fallback: пытаемся найти первые { и последние }
        try:
            start = text.index('{')
            end = text.rindex('}') + 1
            json_str = text[start:end].strip()
            return json.loads(json_str)
        except (ValueError, json.JSONDecodeError) as e:
            logger.error(f"Could not extract JSON from response: {e}")
            logger.debug(f"Response was: {text}")
            raise


    def analyze_input(self, user_text: str, attempt_count: int) -> Dict[str, Any]:
        """
        Анализ ввода пользователя.
        Проверяет длину и является ли текст историей.
        """
        # 1. Проверка длины (менее 50 слов)
        word_count = len(user_text.split())
        if word_count < 50:
            return {
                "status": "insufficient",
                "reason": "length",
                "message": f"Your story is too short ({word_count} words). Please describe the plot in at "
                           f"least 50 words so I can find the best match."
            }

        # 2. Проверка: это история или просто набор слов/вопрос?
        check_prompt = f"""
        Analyze if the following text is a narrative story/plot description or just random words/meta-talk.
        Text: "{user_text}"

        You MUST respond with ONLY valid JSON, nothing else. No explanations before or after.

        Format:
        {{"is_story": true/false, "reason": "brief explanation"}}
        """
        response = ""
        try:
            response = self.llm.complete(check_prompt).text
            logger.debug(f"Story check response: {response[:200]}...")

            # Очистка JSON
            # cleaned_json = re.sub(r"```json|```", "", response).strip()
            # analysis = json.loads(cleaned_json)
            # ✅ Используем улучшенный парсинг
            analysis = self._extract_json(response)


            if not analysis.get("is_story", False):
                return {
                    "status": "insufficient",
                    "reason": "not_story",
                    "message": "This doesn't look like a story. Please describe "
                               "a sequence of events, characters, and what happens to them."
                }
        except Exception as e:
            logger.error(f"Coordinator story check error: {e}")
            logger.debug(f"Full response: {response if 'response' in locals() else 'N/A'}")
            # Fallback: пропускаем, если не удалось проверить
            pass

        # Если все ок
        return {"status": "valid"}

    def generate_suggestion(self, previous_inputs: List[str], genre: str) -> Dict[str, str]:
        """
        Генерирует историю за пользователя, если он не справляется.
        ✅ С валидацией длины (минимум 50 слов)
        """
        # context = " ".join(previous_inputs)
        context = " ".join(previous_inputs[-3:]) if previous_inputs else "nothing specific"

        prompt = f"""
    The user is trying to use a Movie Plot Search engine but fails to provide a good description.
    
    Based on their fragmented inputs: "{context}" (or generate something new if inputs are empty),
    write a compelling, detailed movie plot summary in the {genre} genre.
    
    CRITICAL REQUIREMENTS:
    - MINIMUM 60 words (aim for 70-90 words for a complete plot)
    - Must include: main character(s), conflict, setting, stakes
    - Clear narrative arc with beginning, middle, and potential resolution
    - Engaging and specific details
    - English language only
    
    Genre: {genre}
    
    Output ONLY the story text (60-90 words), no preamble or explanation:
    """

        story_text = self.llm.complete(prompt).text.strip()

        # ✅ ВАЛИДАЦИЯ ДЛИНЫ сгенерированного текста
        word_count = len(story_text.split())

        if word_count < 50:
            logger.warning(f"Generated {genre} plot too short ({word_count} words), expanding...")

            # Повторная генерация с явным требованием расширения
            expansion_prompt = f"""
        The following {genre} plot is TOO SHORT ({word_count} words). 
        Expand it to AT LEAST 60 words while keeping the same theme and characters.
        Add more specific details about the conflict, character motivations, and stakes.

        Current plot: {story_text}

        Expanded version (60-90 words):
        """

            expansion_response = self.llm.complete(expansion_prompt)
            story_text = expansion_response.text.strip()

            # Проверка после расширения
            final_word_count = len(story_text.split())
            logger.info(f"Expanded plot to {final_word_count} words")
        else:
            logger.info(f"Generated {genre} plot has {word_count} words (valid)")

        # Сообщения в зависимости от жанра
        if genre == "romantic":
            msg = ("I see you're having trouble. "
                   " How about we search for a movie with a Romantic plot based on what you said?")
        elif genre == "humorous":
            msg = "Okay, maybe a Humorous story would be better?"
        else:
            msg = f"Let me suggest a {genre} plot for you."

        return {
            "status": "suggestion",
            "genre": genre,
            "message": msg,
            "suggested_story": story_text
        }