



from langgraph.graph import StateGraph, END
from typing import TypedDict, List, Optional, Literal
from datetime import datetime
import json

# Define the state structure for the workflow
class RMIntegrationState(TypedDict):
    current_phase: str
    phase_number: int
    checklist_items: List[dict]
    completed_items: List[int]
    pending_approvals: List[str]
    records_officer_involved: bool
    project_status: str
    comments: dict
    validation_results: dict
    next_phase_ready: bool

# Phase 1: Concept Development
def phase1_concept_development(state: RMIntegrationState) -> RMIntegrationState:
    """Phase 1: Concept Development - Initial records planning"""
    
    phase1_checklist = [
        {
            "id": 1,
            "task": "Include Records Officer in system design process",
            "status": "pending",
            "requires_approval": True,
            "approver": "Records Officer"
        },
        {
            "id": 2,
            "task": "Identify records that support the business process",
            "status": "pending",
            "requires_approval": False,
            "approver": None
        },
        {
            "id": 3,
            "task": "Evaluate current record schedules applicability",
            "status": "pending",
            "requires_approval": False,
            "approver": None
        },
        {
            "id": 4,
            "task": "Determine if new record schedule is required",
            "status": "pending",
            "requires_approval": False,
            "approver": None
        },
        {
            "id": 5,
            "task": "Obtain Records Officer signature on Investment Summary Proposal",
            "status": "pending",
            "requires_approval": True,
            "approver": "Records Officer"
        }
    ]
    
    state["current_phase"] = "Concept Development"
    state["phase_number"] = 1
    state["checklist_items"] = phase1_checklist
    state["pending_approvals"] = ["Records Officer - System Design", "Records Officer - Investment Summary"]
    
    return state

# Phase 2: Requirements Document
def phase2_requirements_document(state: RMIntegrationState) -> RMIntegrationState:
    """Phase 2: Requirements Document - Document all records requirements"""
    
    phase2_checklist = [
        {
            "id": 6,
            "task": "Identify and incorporate all records-related requirements into CONOPS Report",
            "status": "pending",
            "requires_approval": False,
            "approver": None
        },
        {
            "id": 7,
            "task": "Draft new records schedules if needed",
            "status": "pending",
            "requires_approval": False,
            "approver": None
        },
        {
            "id": 8,
            "task": "Obtain Records Officer signature on requirements document",
            "status": "pending",
            "requires_approval": True,
            "approver": "Records Officer"
        }
    ]
    
    state["current_phase"] = "Requirements Document"
    state["phase_number"] = 2
    state["checklist_items"] = phase2_checklist
    state["pending_approvals"] = ["Records Officer - Requirements Document"]
    
    return state

# Phase 3: Design
def phase3_design(state: RMIntegrationState) -> RMIntegrationState:
    """Phase 3: Design - Incorporate records management into system design"""
    
    phase3_checklist = [
        {
            "id": 9,
            "task": "Incorporate records management requirements into system design",
            "status": "pending",
            "requires_approval": False,
            "approver": None
        },
        {
            "id": 10,
            "task": "Obtain Records Officer signature on system design document",
            "status": "pending",
            "requires_approval": True,
            "approver": "Records Officer"
        }
    ]
    
    state["current_phase"] = "Design"
    state["phase_number"] = 3
    state["checklist_items"] = phase3_checklist
    state["pending_approvals"] = ["Records Officer - System Design"]
    
    return state

# Phase 4: Detailed Design
def phase4_detailed_design(state: RMIntegrationState) -> RMIntegrationState:
    """Phase 4: Detailed Design - Include records staff in project meetings"""
    
    phase4_checklist = [
        {
            "id": 11,
            "task": "Include agency records management staff in project status meetings",
            "status": "pending",
            "requires_approval": False,
            "approver": None
        }
    ]
    
    state["current_phase"] = "Detailed Design"
    state["phase_number"] = 4
    state["checklist_items"] = phase4_checklist
    state["pending_approvals"] = []
    
    return state

# Phase 5: Development
def phase5_development(state: RMIntegrationState) -> RMIntegrationState:
    """Phase 5: Development - Continue records staff involvement and submit schedules"""
    
    phase5_checklist = [
        {
            "id": 12,
            "task": "Continue including records management staff in project meetings",
            "status": "pending",
            "requires_approval": False,
            "approver": None
        },
        {
            "id": 13,
            "task": "Submit proposed records schedules to NARA",
            "status": "pending",
            "requires_approval": False,
            "approver": None
        }
    ]
    
    state["current_phase"] = "Development"
    state["phase_number"] = 5
    state["checklist_items"] = phase5_checklist
    state["pending_approvals"] = []
    
    return state

# Phase 6: Integration & System Testing
def phase6_integration_testing(state: RMIntegrationState) -> RMIntegrationState:
    """Phase 6: Integration & System Testing - Test records management integration"""
    
    phase6_checklist = [
        {
            "id": 14,
            "task": "Incorporate records management requirements into system testing",
            "status": "pending",
            "requires_approval": False,
            "approver": None
        },
        {
            "id": 15,
            "task": "Obtain Records Officer signature on Systems Test Report",
            "status": "pending",
            "requires_approval": True,
            "approver": "Records Officer"
        }
    ]
    
    state["current_phase"] = "Integration & System Testing"
    state["phase_number"] = 6
    state["checklist_items"] = phase6_checklist
    state["pending_approvals"] = ["Records Officer - Systems Test Report"]
    
    return state

# Phase 7: Deployment & Acceptance
def phase7_deployment_acceptance(state: RMIntegrationState) -> RMIntegrationState:
    """Phase 7: Deployment & Acceptance - Final approvals and deployment"""
    
    phase7_checklist = [
        {
            "id": 16,
            "task": "Continue including records management staff in project meetings",
            "status": "pending",
            "requires_approval": False,
            "approver": None
        },
        {
            "id": 17,
            "task": "Obtain Records Officer signature on deployment approval document",
            "status": "pending",
            "requires_approval": True,
            "approver": "Records Officer"
        }
    ]
    
    state["current_phase"] = "Deployment & Acceptance"
    state["phase_number"] = 7
    state["checklist_items"] = phase7_checklist
    state["pending_approvals"] = ["Records Officer - Deployment Approval"]
    
    return state

# Phase 8: Production
def phase8_production(state: RMIntegrationState) -> RMIntegrationState:
    """Phase 8: Production - Post-deployment monitoring and compliance"""
    
    phase8_checklist = [
        {
            "id": 18,
            "task": "Complete Mid-Cycle Review (3 years after production)",
            "status": "pending",
            "requires_approval": False,
            "approver": None
        },
        {
            "id": 19,
            "task": "Implement disposition authorities per approved dispositions",
            "status": "pending",
            "requires_approval": False,
            "approver": None
        },
        {
            "id": 20,
            "task": "Send Mid-Cycle Review report to records management staff",
            "status": "pending",
            "requires_approval": False,
            "approver": None
        },
        {
            "id": 21,
            "task": "Obtain Records Officer signature on Mid-Cycle Review certification",
            "status": "pending",
            "requires_approval": True,
            "approver": "Records Officer"
        }
    ]
    
    state["current_phase"] = "Production"
    state["phase_number"] = 8
    state["checklist_items"] = phase8_checklist
    state["pending_approvals"] = ["Records Officer - Mid-Cycle Review Certification"]
    
    return state

# Validation and routing functions
def validate_phase_completion(state: RMIntegrationState) -> RMIntegrationState:
    """Validate that all required items in current phase are completed"""
    
    validation_results = {}
    all_completed = True
    
    for item in state["checklist_items"]:
        if item["status"] != "completed":
            all_completed = False
            validation_results[item["id"]] = f"Item {item['id']} not completed: {item['task']}"
    
    if state["pending_approvals"]:
        all_completed = False
        validation_results["approvals"] = f"Pending approvals: {', '.join(state['pending_approvals'])}"
    
    state["validation_results"] = validation_results
    state["next_phase_ready"] = all_completed
    
    return state

def route_next_phase(state: RMIntegrationState) -> Literal["next_phase", "current_phase", "end"]:
    """Route to next phase or stay in current phase based on completion status"""
    
    if not state["next_phase_ready"]:
        return "current_phase"
    elif state["phase_number"] >= 8:
        return "end"
    else:
        return "next_phase"

def progress_to_next_phase(state: RMIntegrationState) -> RMIntegrationState:
    """Progress to the next phase in the workflow"""
    
    next_phase = state["phase_number"] + 1
    
    # Map phase numbers to phase functions
    phase_functions = {
        1: phase1_concept_development,
        2: phase2_requirements_document,
        3: phase3_design,
        4: phase4_detailed_design,
        5: phase5_development,
        6: phase6_integration_testing,
        7: phase7_deployment_acceptance,
        8: phase8_production
    }
    
    if next_phase in phase_functions:
        return phase_functions[next_phase](state)
    else:
        state["current_phase"] = "Completed"
        state["project_status"] = "All phases completed successfully"
        return state

def stay_current_phase(state: RMIntegrationState) -> RMIntegrationState:
    """Stay in current phase until requirements are met"""
    
    state["project_status"] = f"Phase {state['phase_number']} ({state['current_phase']}) - Requirements not met"
    return state

# Create the workflow graph
def create_rm_integration_workflow():
    """Create the LangGraph workflow for RM Integration"""
    
    workflow = StateGraph(RMIntegrationState)
    
    # Add nodes
    workflow.add_node("phase1", phase1_concept_development)
    workflow.add_node("phase2", phase2_requirements_document)
    workflow.add_node("phase3", phase3_design)
    workflow.add_node("phase4", phase4_detailed_design)
    workflow.add_node("phase5", phase5_development)
    workflow.add_node("phase6", phase6_integration_testing)
    workflow.add_node("phase7", phase7_deployment_acceptance)
    workflow.add_node("phase8", phase8_production)
    workflow.add_node("validate", validate_phase_completion)
    workflow.add_node("next_phase", progress_to_next_phase)
    workflow.add_node("current_phase", stay_current_phase)
    
    # Set entry point
    workflow.set_entry_point("phase1")
    
    # Add edges
    workflow.add_edge("phase1", "validate")
    workflow.add_edge("phase2", "validate")
    workflow.add_edge("phase3", "validate")
    workflow.add_edge("phase4", "validate")
    workflow.add_edge("phase5", "validate")
    workflow.add_edge("phase6", "validate")
    workflow.add_edge("phase7", "validate")
    workflow.add_edge("phase8", "validate")
    
    # Add conditional edges from validate
    workflow.add_conditional_edges(
        "validate",
        route_next_phase,
        {
            "next_phase": "next_phase",
            "current_phase": "current_phase", 
            "end": END
        }
    )
    
    # Add edge from next_phase back to appropriate phase
    workflow.add_edge("next_phase", "validate")
    workflow.add_edge("current_phase", END)
    
    return workflow.compile()

# Example usage and testing
def run_rm_integration_example():
    """Example of how to run the RM integration workflow"""
    
    # Initialize the workflow
    app = create_rm_integration_workflow()
    
    # Initial state
    initial_state = {
        "current_phase": "",
        "phase_number": 0,
        "checklist_items": [],
        "completed_items": [],
        "pending_approvals": [],
        "records_officer_involved": False,
        "project_status": "Starting RM Integration Process",
        "comments": {},
        "validation_results": {},
        "next_phase_ready": False
    }
    
    # Run the workflow
    result = app.invoke(initial_state)
    
    return result

# Utility functions for workflow management
def update_item_status(state: RMIntegrationState, item_id: int, status: str, comment: str = "") -> RMIntegrationState:
    """Update the status of a specific checklist item"""
    
    for item in state["checklist_items"]:
        if item["id"] == item_id:
            item["status"] = status
            if status == "completed" and item_id not in state["completed_items"]:
                state["completed_items"].append(item_id)
            break
    
    if comment:
        state["comments"][item_id] = comment
    
    return state

def generate_status_report(state: RMIntegrationState) -> dict:
    """Generate a comprehensive status report"""
    
    completed_count = len([item for item in state["checklist_items"] if item["status"] == "completed"])
    total_count = len(state["checklist_items"])
    
    report = {
        "current_phase": state["current_phase"],
        "phase_number": state["phase_number"],
        "completion_percentage": (completed_count / total_count * 100) if total_count > 0 else 0,
        "completed_items": completed_count,
        "total_items": total_count,
        "pending_approvals": state["pending_approvals"],
        "validation_results": state["validation_results"],
        "project_status": state["project_status"],
        "timestamp": datetime.now().isoformat()
    }
    
    return report

def test_ollama_similarity(requirement: str) -> Dict[str, Any]:
    """
    Call the /test_ollama_similarity endpoint with a requirement text
    
    Args:
        requirement: The requirement text to generate embeddings for
        
    Returns:
        Dictionary containing embedding preview, dimensions, and model info
    """
    try:
        # Prepare the request payload
        payload = {
            "prompt": requirement
        }
        
        # Make POST request to the API
        response = requests.post(
            f"{API_BASE_URL}/test_ollama_similarity",
            json=payload,
            headers={"Content-Type": "application/json"}
        )
        
        # Check if request was successful
        response.raise_for_status()
        
        # Parse and return the JSON response
        result = response.json()
        
        return result
        
    except requests.exceptions.RequestException as e:
        return {"error": f"Request failed: {str(e)}"}
    except json.JSONDecodeError as e:
        return {"error": f"Failed to parse response: {str(e)}"}

def test_multiple_requirements(requirements: list) -> list:
    """
    Process multiple requirements and return their embeddings
    
    Args:
        requirements: List of requirement texts
        
    Returns:
        List of results for each requirement
    """
    results = []
    
    for i, req in enumerate(requirements):
        result = test_ollama_similarity(req)
        results.append({
            "requirement": req,
            "result": result
        })
    
    return results

def compare_requirements_similarity(req1: str, req2: str) -> Dict[str, Any]:
    """
    Compare similarity between two requirements using cosine similarity
    
    Args:
        req1: First requirement text
        req2: Second requirement text
        
    Returns:
        Dictionary with embeddings and similarity score
    """
    import numpy as np
    
    # Get embeddings for both requirements
    result1 = test_ollama_similarity(req1)
    result2 = test_ollama_similarity(req2)
    
    if "error" in result1 or "error" in result2:
        return {
            "error": "Failed to generate embeddings",
            "result1": result1,
            "result2": result2
        }
    
    # For full comparison, we need the complete embeddings
    # This is just a demo with preview data
    emb1 = np.array(result1.get('embedding_preview', []))
    emb2 = np.array(result2.get('embedding_preview', []))
    
    # Calculate cosine similarity
    if len(emb1) > 0 and len(emb2) > 0:
        dot_product = np.dot(emb1, emb2)
        norm1 = np.linalg.norm(emb1)
        norm2 = np.linalg.norm(emb2)
        similarity = dot_product / (norm1 * norm2)
    else:
        similarity = None
    
    return {
        "requirement1": req1,
        "requirement2": req2,
        "similarity_score": float(similarity) if similarity is not None else None
    }



if __name__ == "__main__":
    # Run example
    result = run_rm_integration_example()
    print("RM Integration Workflow Result:")
    print(json.dumps(generate_status_report(result), indent=2))