# agents.py
from typing import Dict, Any
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage, SystemMessage
from langchain_community.chat_models import ChatOllama
import json
from config import MODEL_NAME, TEMPERATURE
from models import CodeSolution, ReviewResult, ReviewComment, Task

class DeveloperAgent:
    """Агент-разработчик: генерирует код по требованию"""
    
    def __init__(self):
        self.llm = ChatOllama(
            model=MODEL_NAME,
            temperature=TEMPERATURE,
            base_url="http://localhost:11434"
        )
    
    def develop(self, task: Task, previous_comments: list = None) -> CodeSolution:
        """Генерирует код на основе требования и замечаний"""
        
        system_prompt = """Ты - опытный разработчик. Твоя задача - написать качественный код.
        Ты должен:
        1. Писать чистый, читаемый код с комментариями
        2. Добавлять обработку ошибок
        3. Писать юнит-тесты (pytest/unittest)
        4. Следовать PEP8 (для Python) или стандартам языка
        5. Возвращать решение в строго определённом JSON формате
        
        В ответе должен быть JSON с полями:
        - files: dict (имя_файла -> содержимое)
        - description: str (описание решения)
        - tests: dict (имя_теста -> содержимое теста)
        """
        
        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)
        
        # Парсим JSON ответ
        try:
            # Извлекаем JSON из ответа
            content = response.content
            # Находим JSON в тексте (между ```json и ``` или просто сам JSON)
            if "```json" in content:
                content = content.split("```json")[1].split("```")[0]
            elif "```" in content:
                content = content.split("```")[1].split("```")[0]
            
            data = json.loads(content.strip())
            return CodeSolution(
                files=data.get("files", {}),
                description=data.get("description", ""),
                tests=data.get("tests", {})
            )
        except Exception as e:
            raise ValueError(f"Failed to parse developer response: {e}\nResponse: {response.content}")

class ReviewerAgent:
    """Агент-ревьювер: проверяет качество кода"""
    
    def __init__(self):
        self.llm = ChatOllama(
            model=MODEL_NAME,
            temperature=0.1,
            base_url="http://localhost:11434"
        )
    
    def review(self, task: Task, solution: CodeSolution) -> ReviewResult:
        """Проверяет код и возвращает замечания"""
        
        system_prompt = """Ты - строгий ревьювер кода. Твоя задача:
        1. Проверить соответствие требованиям
        2. Найти логические ошибки и баги
        3. Оценить читаемость и стиль кода
        4. Проверить покрытие тестами
        5. Найти проблемы безопасности
        6. Предложить улучшения
        
        Оценивай код критически. Будь конкретен в замечаниях, указывай файлы и строки.
        
        Верни JSON с полями:
        - status: "approved", "changes_requested", или "rejected"
        - comments: массив объектов с полями: file, line, severity, text
        - summary: краткое резюме ревью
        """
        
        # Формируем представление кода для ревью
        code_for_review = ""
        for filename, content in solution.files.items():
            code_for_review += f"\n--- {filename} ---\n{content}\n"
        
        if solution.tests:
            code_for_review += "\n--- TESTS ---\n"
            for testname, test_content in solution.tests.items():
                code_for_review += f"\n--- {testname} ---\n{test_content}\n"
        
        user_prompt = f"""
        Требование: {task.requirement}
        
        Код на ревью:
        {code_for_review}
        
        Описание решения разработчика: {solution.description}
        
        Проведи ревью. Будь строг, но справедлив.
        """
        
        messages = [
            SystemMessage(content=system_prompt),
            HumanMessage(content=user_prompt)
        ]
        
        response = self.llm.invoke(messages)
        
        # Парсим JSON ответ
        try:
            content = response.content
            if "```json" in content:
                content = content.split("```json")[1].split("```")[0]
            elif "```" in content:
                content = content.split("```")[1].split("```")[0]
            
            data = json.loads(content.strip())
            
            comments = []
            for comment_data in data.get("comments", []):
                comments.append(ReviewComment(
                    line=comment_data.get("line"),
                    file=comment_data.get("file"),
                    severity=comment_data.get("severity", "major"),
                    text=comment_data.get("text", "")
                ))
            
            return ReviewResult(
                status=data.get("status", "changes_requested"),
                comments=comments,
                summary=data.get("summary", "")
            )
        except Exception as e:
            raise ValueError(f"Failed to parse reviewer response: {e}\nResponse: {response.content}")