import os
import time
from typing import TypedDict, Literal
from langgraph.graph import StateGraph, END
from pydantic import BaseModel, Field

# Define the structured output for the Critic
class ReviewResponse(BaseModel):
    status: Literal["approved", "rejected"] = Field(
        description="Whether the document meets all criteria or needs revision."
    )
    feedback: str = Field(
        description="Detailed feedback on what is missing or needs improvement. If approved, provide a brief summary of why."
    )

# Define the State
class WorkflowState(TypedDict):
    raw_input: str
    current_draft: str
    review_feedback: str
    status: str
    revision_count: int

# Nodes
def writer_node(state: WorkflowState):
    # Simulate thinking time
    time.sleep(2)
    
    raw_input = state.get("raw_input", "")
    revision_count = state.get("revision_count", 0)

    # Mock PRD generation based on input
    mock_prd = f"""# Product Requirement Document (PRD)

## Executive Summary
This document outlines the requirements based on the raw notes: "{raw_input[:50]}..."

## User Personas
- **Primary User**: General users looking for seamless interaction.
- **Admin**: System administrators managing access.

## User Stories
1. As a user, I want to log in securely so that I can access my dashboard.
2. As an admin, I want to view active sessions.

## Technical Specifications
- **Frontend**: HTML/CSS/JS (Vanilla & Glassmorphism)
- **Backend**: FastAPI with WebSockets
- **Architecture**: Multi-Agent LangGraph

## Edge Cases
- Invalid credentials gracefully handled.
- Rate limits triggered and warnings shown.

## Out of Scope
- Third-party SSO (except where specified).
- Mobile App development.

> **Note**: This is a mock PRD generated locally without an OpenAI API key. Revision count: {revision_count + 1}
"""
    
    return {
        "current_draft": mock_prd,
        "revision_count": revision_count + 1
    }

def critic_node(state: WorkflowState):
    # Simulate thinking time
    time.sleep(2)
    
    revision_count = state.get("revision_count", 0)
    
    # Mock Review Logic: Reject the first time to show the loop, approve the second time.
    if revision_count < 2:
        return {
            "status": "rejected",
            "review_feedback": "The Technical Specifications section needs more detail regarding the database schema. Please revise."
        }
    else:
        return {
            "status": "approved",
            "review_feedback": "The PRD looks excellent and meets all company standards. Moving to human approval."
        }

def human_approval_node(state: WorkflowState):
    # This node doesn't actually do anything automatically.
    # It serves as a placeholder for the graph to pause.
    # When the human approves via API, we will just pass the state through.
    return {"status": "human_approved"}

# Routing
def route_after_critic(state: WorkflowState):
    if state["status"] == "approved":
        return "human_approval"
    elif state["revision_count"] >= 3:
        # Prevent infinite loops, force human review if too many revisions
        return "human_approval"
    else:
        return "writer"

def build_workflow():
    workflow = StateGraph(WorkflowState)
    
    # Add nodes
    workflow.add_node("writer", writer_node)
    workflow.add_node("critic", critic_node)
    workflow.add_node("human_approval", human_approval_node)
    
    # Set entry point
    workflow.set_entry_point("writer")
    
    # Add edges
    workflow.add_edge("writer", "critic")
    workflow.add_conditional_edges(
        "critic",
        route_after_critic,
        {
            "writer": "writer",
            "human_approval": "human_approval"
        }
    )
    workflow.add_edge("human_approval", END)
    
    # Compile graph with a checkpointer for human-in-the-loop (memory)
    from langgraph.checkpoint.memory import MemorySaver
    memory = MemorySaver()
    
    # The graph will interrupt BEFORE the human_approval node
    app = workflow.compile(
        checkpointer=memory,
        interrupt_before=["human_approval"]
    )
    
    return app

# Singleton instance of the graph
workflow_app = build_workflow()