from typing import Dict, Any, List
from typing_extensions import TypedDict

class AgentState(TypedDict):
    """
    AgentState extends HiveGPTMemoryState to provide a unified state structure for the Compute Agent workflow.
    
    Inherits all memory-related fields and adds compute agent-specific workflow fields:
    
    Core Fields:
        - query: User's input query
        - response: Final response to user
        - current_step: Current workflow step identifier
        
    Decision Fields:
        - agent_decision: Routing decision ('deploy_model' or 'react_agent')
        - deployment_approved: Whether human approved deployment
        
    Model Deployment Fields:
        - model_name: Name/ID of the model to deploy
        - model_card: Raw model card data from HuggingFace
        - model_info: Extracted model information (JSON)
        - capacity_estimate: Estimated compute resources needed
        - deployment_result: Result of model deployment
        
    React Agent Fields:
        - react_results: Results from React agent execution
        - tool_calls: List of tool calls made by React agent
        - tool_results: Results from tool executions
        
    Error Handling:
        - error: Error message if any step fails
        - error_step: Step where error occurred
    """
    # Core fields
    query: str
    response: str
    current_step: str
    messages: List[Dict[str, Any]]
    
    # Decision fields
    agent_decision: str
    deployment_approved: bool
    
    # Model deployment fields
    model_name: str
    model_card: Dict[str, Any]
    model_info: Dict[str, Any]
    capacity_estimate: Dict[str, Any]
    deployment_result: Dict[str, Any]
    capacity_estimation_status: str
    capacity_approval_status: str
    capacity_approved: bool
    estimated_gpu_memory: float
    gpu_requirements: Dict[str, Any]
    cost_estimates: Dict[str, Any]
    need_reestimation: bool
    
    # React agent fields
    react_results: Dict[str, Any]
    tool_calls: List[Dict[str, Any]]
    tool_results: List[Dict[str, Any]]

    # Tool approval fields (for human-in-the-loop)
    pending_tool_calls: List[Dict[str, Any]]
    approved_tool_calls: List[Dict[str, Any]]
    rejected_tool_calls: List[Dict[str, Any]]
    modified_tool_calls: List[Dict[str, Any]]
    needs_re_reasoning: bool
    re_reasoning_feedback: str

    # User identification
    user_id: str
    session_id: str

    # Workflow identification (for tools registry lookup)
    workflow_id: int

    # Compute instance fields (for HiveCompute deployment)
    instance_id: str
    instance_status: str
    instance_created: bool