from typing import Optional, Dict, Any
from langchain_core.messages import HumanMessage, SystemMessage
from langchain_ollama import ChatOllama
import json
from config import MODEL_NAME, TEMPERATURE, LLM_HOST
from models import CodeSolution, ReviewResult, ReviewComment, Task
import re

class DeveloperAgent:
    """Агент-разработчик: генерирует код по требованию"""
    
    def __init__(self):
        self.llm = ChatOllama(
            model=MODEL_NAME,
            temperature=TEMPERATURE,
            base_url=LLM_HOST,
            num_predict=8192
        )
    
    def develop(self, task: Task, previous_comments: list = None) -> CodeSolution:
        """Генерирует код с улучшенным парсингом JSON"""
        
        system_prompt = """Ты - опытный разработчик. Твоя задача - написать качественный код.
        Ты должен:
        1. Писать чистый, читаемый код с комментариями
        2. Добавлять обработку ошибок
        3. Писать юнит-тесты (pytest/unittest)
        4. Следовать PEP8 (для Python) или стандартам языка
        
        ВАЖНО: Возвращай ТОЛЬКО валидный JSON без лишних символов.
        ЗАПРЕЩЕНЫ trailing commas (запятые перед закрывающей скобкой).
        
        Формат ответа (пример):
        {
            "files": {
                "main.py": "print('hello')"
            },
            "description": "Краткое описание",
            "tests": {
                "test_main.py": "def test_hello(): assert True"
            }
        }
        
        Не добавляй пояснений вне JSON.
        """
        
        user_prompt = f"""
        Требование: {task.requirement}
        Текущая итерация: {task.iteration}
        """
        
        if previous_comments:
            user_prompt += f"""
            
            Замечания с предыдущего ревью:
            {json.dumps(previous_comments, indent=2, ensure_ascii=False)}
            
            Пожалуйста, исправь код согласно замечаниям.
            """
        
        messages = [
            SystemMessage(content=system_prompt),
            HumanMessage(content=user_prompt)
        ]
        
        response = self.llm.invoke(messages)
        raw_content = response.content
        
        # Пытаемся извлечь JSON из разных форматов
        parsed_data = self._extract_json(raw_content)
        
        if parsed_data is None:
            # Если не удалось, пробуем повторно с explicit инструкцией
            print("⚠️ Не удалось распарсить JSON, повторный запрос с исправлением...")
            correction_prompt = f"""
            Предыдущий ответ содержал невалидный JSON (возможно, trailing commas).
            Пожалуйста, исправь и верни ТОЛЬКО валидный JSON в формате:
            {{
                "files": {{"filename": "content"}},
                "description": "text",
                "tests": {{"testfile": "test_content"}}
            }}
            
            Твоя предыдущая попытка:
            {raw_content[:1000]}
            
            Требование: {task.requirement}
            """
            
            messages2 = [
                SystemMessage(content="Ты - разработчик. Верни только валидный JSON без лишних символов."),
                HumanMessage(content=correction_prompt)
            ]
            response2 = self.llm.invoke(messages2)
            parsed_data = self._extract_json(response2.content)
            
            if parsed_data is None:
                raise ValueError(f"Failed to parse JSON after retry. Raw: {response2.content[:500]}")
        
        return CodeSolution(
            files=parsed_data.get("files", {}),
            description=parsed_data.get("description", ""),
            tests=parsed_data.get("tests", {})
        )
    
    def _extract_json(self, text: str) -> Optional[dict]:
        """Извлекает и очищает JSON из текста ответа LLM"""
        
        # 1. Извлечь блок между ```json и ``` или просто ```
        if "```json" in text:
            text = text.split("```json")[1].split("```")[0]
        elif "```" in text:
            text = text.split("```")[1].split("```")[0]
        
        text = text.strip()
        
        # 2. Удалить trailing commas в объектах и массивах
        # Удаляем запятую перед } или ]
        text = re.sub(r',\s*}', '}', text)
        text = re.sub(r',\s*]', ']', text)
        
        # 3. Удалить комментарии (// ...) если есть
        text = re.sub(r'//.*?(\n|$)', '', text)
        
        # 4. Попробовать распарсить
        try:
            return json.loads(text)
        except json.JSONDecodeError as e:
            print(f"JSON decode error: {e}")
            # 5. Еще одна попытка: найти любой JSON объект в строке
            match = re.search(r'\{.*\}|\[.*\]', text, re.DOTALL)
            if match:
                try:
                    return json.loads(match.group())
                except:
                    pass
            return None