# orchestrator.py
from typing import Dict, Any, Literal, TypedDict, List, Optional
from langgraph.graph import StateGraph, END
from langgraph.checkpoint.memory import MemorySaver
import uuid
from datetime import datetime
from models import Task, TaskStatus, CodeSolution, ReviewResult, IterationRecord
from agent.developer import DeveloperAgent
from agent.reviewer import ReviewerAgent
from config import MAX_REVIEW_RETRIES

class AgentState(TypedDict):
    task: Task
    plan: Dict
    execution_results: List[Dict]
    review: Optional[ReviewResult]
    iteration: int
    max_iterations: int

class Orchestrator:
    """Оркестратор - управляет процессом разработки"""
    
    def __init__(self):
        self.developer = DeveloperAgent()
        self.reviewer = ReviewerAgent()
        self.max_retries = MAX_REVIEW_RETRIES
        
        # Создаем граф состояний
        self.workflow = self._build_workflow()
        self.checkpointer = MemorySaver()
        
    def _add_history(self, task: Task, agent: str, action: str, 
                    input_summary: str = None, output_summary: str = None,
                    details: Dict = None, error: str = None):
        """Добавляет запись в историю задачи"""
        record = IterationRecord(
            iteration=task.iteration,
            timestamp=datetime.now().isoformat(),
            agent=agent,
            action=action,
            input_summary=input_summary,
            output_summary=output_summary,
            details=details or {},
            error=error
        )
        task.history.append(record)
    
    def _build_workflow(self):
        """Строит граф процесса разработки"""
        
        # Определяем состояния
        workflow = StateGraph(AgentState)
        
        # Добавляем узлы
        workflow.add_node("planner", self._plan_node)
        workflow.add_node("executor", self._exec_node)
        workflow.add_node("reviewer", self._review_node)
        
        # Определяем переходы
        workflow.set_entry_point("planner")
        workflow.add_edge("planner", "executor")
        workflow.add_edge("executor", "reviewer")
        
        workflow.add_conditional_edges("reviewer", self._after_review, {
            "approved": END,
            "rework": "executor",  # повторить выполнение с тем же планом? или перепланировать?
            "replan": "planner",
            "reject": END
        })
        
        return workflow.compile()
    
    def _plan_node(self, state: AgentState):
        task = state["task"]
        print(f"\n🔧 [Iteration {task.iteration}] Планировщик начал работу...")
        
        # Создаем запись в истории
        task.history.append(IterationRecord(
            iteration=task.iteration,
            timestamp=datetime.now().isoformat(),
            agent="developer",
            action="start",
            input_summary=f"Requirement: {task.requirement[:100]}...",
            details={}
        ))
        
        # План просто содержит задачу - разработчик выполнит её напрямую
        plan = {"task": task.requirement, "iteration": task.iteration}
        state["plan"] = plan
        return state

    def _exec_node(self, state: AgentState):
        """Узел разработчика"""
        task = state["task"]
        print(f"\n🔧 [Iteration {task.iteration}] Исполнитель начал работу...")
        
        # Обновляем статус
        task.status = TaskStatus.DEVELOPMENT
        
        # Разработчик выполняет задачу напрямую
        solution = self.developer.develop(task, [])
        task.solution = solution
        task.status = TaskStatus.REVIEW
        
        # Записываем результат выполнения
        state["execution_results"] = [{"status": "success", "solution": solution}]
        
        # Добавляем запись в историю
        task.history.append(IterationRecord(
            iteration=task.iteration,
            timestamp=datetime.now().isoformat(),
            agent="developer",
            action="complete",
            output_summary=f"Generated {len(solution.files)} file(s), {len(solution.tests)} test(s)",
            details={
                "files": list(solution.files.keys()),
                "tests": list(solution.tests.keys())
            }
        ))
        
        print(f"✅ Разработчик создал решение: {len(solution.files)} файлов")
        return state

    def _review_node(self, state: AgentState):
        task = state["task"]
        print(f"\n🔍 [Iteration {task.iteration}] Ревьювер начал работу...")
        
        # Добавляем запись о начале ревью
        task.history.append(IterationRecord(
            iteration=task.iteration,
            timestamp=datetime.now().isoformat(),
            agent="reviewer",
            action="start",
            input_summary=f"Reviewing {len(task.solution.files) if task.solution else 0} file(s)",
            details={}
        ))
        
        if not task.solution:
            task.status = TaskStatus.FAILED
            error_msg = "No solution to review"
            task.error_log.append(error_msg)
            task.history.append(IterationRecord(
                iteration=task.iteration,
                timestamp=datetime.now().isoformat(),
                agent="reviewer",
                action="error",
                error=error_msg
            ))
            return state
        
        # Проводим ревью
        review = self.reviewer.review(task, task.solution)
        task.last_review = review
        task.iteration += 1
        
        # Добавляем запись в историю
        task.history.append(IterationRecord(
            iteration=task.iteration - 1,  # предыдущая итерация
            timestamp=datetime.now().isoformat(),
            agent="reviewer",
            action="complete",
            output_summary=f"Verdict: {review.status}, Comments: {len(review.comments)}",
            details={
                "verdict": review.status,
                "comments_count": len(review.comments)
            }
        ))
        
        # Сохраняем результат ревью
        state["review"] = review
        
        print(f"📊 Вердикт: {review.status}")
        if review.comments:
            print(f"   Найдено замечаний: {len(review.comments)}")
        
        return state

    def _after_review(self, state: AgentState):
        task = state["task"]
        review = state["review"]
        
        # Определяем статус
        if isinstance(review, ReviewResult):
            status = review.status
        elif isinstance(review, dict):
            status = review.get("status")
        else:
            status = None
        
        if status == "approved":
            task.status = TaskStatus.APPROVED
            return "approved"
        elif status == "changes_requested":
            if task.iteration >= state.get("max_iterations", 3):
                task.status = TaskStatus.FAILED
                return "reject"
            task.status = TaskStatus.REWORK
            return "rework"
        elif status == "rejected":
            if task.iteration >= state.get("max_iterations", 3):
                task.status = TaskStatus.FAILED
                return "reject"
            task.status = TaskStatus.REWORK
            return "replan"
        else:
            task.status = TaskStatus.FAILED
            return "reject"
    
    def _save_solution(self, solution: CodeSolution):
        """Сохраняет решение в файлы"""
        import os
        output_dir = "generated_code"
        os.makedirs(output_dir, exist_ok=True)
        
        for filename, content in solution.files.items():
            filepath = os.path.join(output_dir, filename)
            with open(filepath, "w", encoding="utf-8") as f:
                f.write(content)
            print(f"   💾 Сохранен файл: {filepath}")
        
        for testname, test_content in solution.tests.items():
            testpath = os.path.join(output_dir, testname)
            with open(testpath, "w", encoding="utf-8") as f:
                f.write(test_content)
            print(f"   💾 Сохранен тест: {testpath}")
    
    def run(self, task_text: str, max_iterations: int = 3) -> Dict:
        """Запускает процесс разработки"""
        
        task_id = str(uuid.uuid4())
        
        # Создаем объект Task
        task = Task(
            id=task_id,
            requirement=task_text,
            status=TaskStatus.PENDING,
            iteration=0,
            solution=None,
            last_review=None,
            error_log=[],
            history=[]
        )
        
        initial_state = {
            "task": task,
            "plan": {},
            "execution_results": [],
            "review": None,
            "iteration": 0,
            "max_iterations": max_iterations
        }
        
        print(f"\n🚀 Запуск мультиагентной системы")
        print(f"📝 Задача: {task_text[:100]}...")
        print(f"🆔 ID: {task_id}")
        
        # Запускаем граф
        final_state = self.workflow.invoke(initial_state)
        
        # Сохраняем решение если успешно
        final_task = final_state["task"]
        if final_task.status == TaskStatus.APPROVED and final_task.solution:
            self._save_solution(final_task.solution)
        
        print("\n" + "="*50)
        print("📋 ИТОГИ РАБОТЫ")
        print("="*50)
        print(f"Статус: {final_task.status.value}")
        print(f"Итераций: {final_task.iteration}")
        
        if final_task.status == TaskStatus.APPROVED:
            print(f"\n✅ Решение принято!")
            print(f"   Файлов: {len(final_task.solution.files)}")
            print(f"   Тестов: {len(final_task.solution.tests)}")
        elif final_task.last_review:
            print(f"\n📝 Финальный вердикт: {final_task.last_review.summary}")
        
        if final_task.error_log:
            print(f"\n❌ Ошибки: {len(final_task.error_log)}")
            for error in final_task.error_log:
                print(f"   - {error}")
        
        return {"task": final_task}