from typing import Dict, Any, TypedDict, Annotated, Sequence
from langgraph.graph import Graph, StateGraph
from langgraph.prebuilt import ToolExecutor
from agents.business_analyst import BusinessAnalystAgent
from agents.designer import DesignerAgent
from agents.developer import DeveloperAgent
from agents.tester import TesterAgent
from core.database import db

class AgentState(TypedDict):
    """State for the agent workflow"""
    requirements: str
    user_stories: str
    design: str
    code: str
    test_cases: str
    test_results: str
    current_step: str
    status: str
    messages: Sequence[str]

class DevelopmentWorkflow:
    def __init__(self):
        self.business_analyst = BusinessAnalystAgent()
        self.designer = DesignerAgent()
        self.developer = DeveloperAgent()
        self.tester = TesterAgent()
        
        # Define the workflow
        self.workflow = self._create_workflow()
        
    def _create_workflow(self) -> Graph:
        """Create the LangGraph workflow"""
        workflow = StateGraph(AgentState)
        
        # Add nodes for each step
        workflow.add_node("create_user_stories", self._create_user_stories)
        workflow.add_node("create_design", self._create_design)
        workflow.add_node("generate_code", self._generate_code)
        workflow.add_node("create_tests", self._create_tests)
        workflow.add_node("execute_tests", self._execute_tests)
        
        # Define the edges
        workflow.add_edge("create_user_stories", "create_design")
        workflow.add_edge("create_design", "generate_code")
        workflow.add_edge("generate_code", "create_tests")
        workflow.add_edge("create_tests", "execute_tests")
        
        # Set entry point
        workflow.set_entry_point("create_user_stories")
        
        # Set exit point
        workflow.set_finish_point("execute_tests")
        
        return workflow.compile()
    
    async def _create_user_stories(self, state: AgentState) -> AgentState:
        """Create user stories from requirements"""
        result = await self.business_analyst.create_user_stories(state["requirements"])
        return {
            **state,
            "user_stories": result["user_stories"],
            "current_step": "create_user_stories",
            "status": "in_progress",
            "messages": [*state["messages"], "User stories created successfully"]
        }
    
    async def _create_design(self, state: AgentState) -> AgentState:
        """Create design from user stories"""
        result = await self.designer.create_design(state["user_stories"])
        return {
            **state,
            "design": result["design"],
            "current_step": "create_design",
            "status": "in_progress",
            "messages": [*state["messages"], "Design created successfully"]
        }
    
    async def _generate_code(self, state: AgentState) -> AgentState:
        """Generate code from design"""
        result = await self.developer.generate_code(state["design"])
        return {
            **state,
            "code": result["code"],
            "current_step": "generate_code",
            "status": "in_progress",
            "messages": [*state["messages"], "Code generated successfully"]
        }
    
    async def _create_tests(self, state: AgentState) -> AgentState:
        """Create test cases from code"""
        result = await self.tester.create_test_cases(state["code"])
        return {
            **state,
            "test_cases": result["test_cases"],
            "current_step": "create_tests",
            "status": "in_progress",
            "messages": [*state["messages"], "Test cases created successfully"]
        }
    
    async def _execute_tests(self, state: AgentState) -> AgentState:
        """Execute tests on the code"""
        result = await self.tester.execute_tests(state["code"], state["test_cases"])
        return {
            **state,
            "test_results": result["test_results"],
            "current_step": "execute_tests",
            "status": "completed",
            "messages": [*state["messages"], "Tests executed successfully"]
        }
    
    async def run(self, requirements: str) -> Dict[str, Any]:
        """Run the complete workflow"""
        initial_state = AgentState(
            requirements=requirements,
            user_stories="",
            design="",
            code="",
            test_cases="",
            test_results="",
            current_step="start",
            status="initialized",
            messages=[]
        )
        
        final_state = await self.workflow.ainvoke(initial_state)
        
        return {
            "status": "success",
            "user_stories": final_state["user_stories"],
            "design": final_state["design"],
            "code": final_state["code"],
            "test_results": final_state["test_results"],
            "messages": final_state["messages"]
        }