import json
import os
import asyncio
from google.adk.agents import LlmAgent, SequentialAgent
from google.adk.tools import FunctionTool
from google.adk.runners import Runner
from google.adk.sessions import InMemorySessionService
from google.genai import types
from pydantic import BaseModel, Field
import nest_asyncio
import agentops
from dotenv import load_dotenv

# Load environment variables from .env file
load_dotenv()
nest_asyncio.apply()
AGENTOPS_API_KEY = os.getenv("AGENTOPS_API_KEY") or "your_agentops_api_key_here"
# Initialize AgentOps - Just 2 lines!
agentops.init(AGENTOPS_API_KEY, trace_name="adk-human-approval-notebook", auto_start_session=False)

APP_NAME = "human_approval_app_notebook"
USER_ID = "test_user_notebook_123"
SESSION_ID = "approval_session_notebook_456"
MODEL_NAME = "gemini-1.5-flash"
agentops.start_trace(trace_name=APP_NAME, tags=["google_adk","notebook"])

class ApprovalRequest(BaseModel):
    amount: float = Field(description="The amount requiring approval")
    reason: str = Field(description="The reason for the request")
class ApprovalDecision(BaseModel):
    decision: str = Field(description="The approval decision: 'approved' or 'rejected'")
    comments: str = Field(description="Additional comments from the approver")

async def external_approval_tool(amount: float, reason: str) -> str:
    """ 
    Prompts for human approval and returns the decision as a JSON string.
    """
    print(f"🔔 HUMAN APPROVAL REQUIRED:")
    print(f"   Amount: ${amount:,.2f}")
    print(f"   Reason: {reason}")
    decision = ""
    while decision.lower() not in ["approved", "rejected"]:
        decision = input("   Enter decision (approved/rejected): ").strip().lower()
        if decision.lower() not in ["approved", "rejected"]:
            print("   Invalid input. Please enter 'approved' or 'rejected'.")
    comments = input("   Enter comments (optional): ").strip()
    print(f"   Decision: {decision.upper()}")
    print(f"   Comments: {comments if comments else 'N/A'}")
    return json.dumps({
        "decision": decision,
        "comments": comments,
        "amount": amount,
        "reason": reason
    })

# Create the approval tool instance
approval_tool = FunctionTool(func=external_approval_tool)

# Agent 1: Prepare the approval request
prepare_request = LlmAgent(
    model=MODEL_NAME,
    name="PrepareApprovalAgent",
    description="Extracts and prepares approval request details from user input",
    instruction="""You are an approval request preparation agent.
        Your task:
        1. Extract the amount and reason from the user's request
        2. Store these values in the session state with keys 'approval_amount' and 'approval_reason'
        3. Validate that both amount and reason are provided
        4. Respond with a summary of what will be submitted for approval
    If the user input is missing amount or reason, ask for clarification.
    """,
    output_key="request_prepared"
)

# Agent 2: Request human approval using the tool
request_approval = LlmAgent(
    model=MODEL_NAME,
    name="RequestHumanApprovalAgent", 
    description="Calls the external approval system with prepared request details",
    instruction="""You are a human approval request agent.
        Your task:
        1. Get the 'approval_amount' and 'approval_reason' from the session state
        2. Use the external_approval_tool with these values
        3. Store the approval decision in session state with key 'human_decision'
        4. Respond with the approval status
    Always use the exact values from the session state for the tool call.
    """,
    tools=[approval_tool],
    output_key="approval_requested"
)

# Agent 3: Process the approval decision
process_decision = LlmAgent(
    model=MODEL_NAME,
    name="ProcessDecisionAgent",
    description="Processes the human approval decision and provides final response",
    instruction="""You are a decision processing agent.
        Your task:
        1. Check the 'human_decision' from session state
        2. Parse the approval decision JSON
        3. If approved: congratulate and provide next steps
        4. If rejected: explain the rejection and suggest alternatives
        5. Provide a clear, helpful final response to the user

    Be professional and helpful in your response.
    """,
    output_key="final_decision"
)

approval_workflow = SequentialAgent(
    name="HumanApprovalWorkflowNotebook",
    description="Complete workflow for processing approval requests with human oversight",
    sub_agents=[prepare_request, request_approval, process_decision]
)

session_service = InMemorySessionService()
# Create runner
workflow_runner = Runner(
    agent=approval_workflow,
    app_name=APP_NAME,
    session_service=session_service
)

async def run_approval_workflow_notebook(user_request: str, session_id: str):
    """Run the complete approval workflow with a user request in the notebook environment"""
    print(f"{'='*60}")
    print(f" Starting Approval Workflow for Session: {session_id}")
    print(f"{'='*60}")
    print(f"User Request: {user_request}")
    # Create user message
    user_content = types.Content(
        role='user', 
        parts=[types.Part(text=user_request)]
    )
    step_count = 0
    final_response = "No response received"
    # Run the workflow
    async for event in workflow_runner.run_async(
        user_id=USER_ID,
        session_id=session_id,
        new_message=user_content,
    ):
        if event.author and event.content:
            step_count += 1
            print(f"📋 Step {step_count} - {event.author}:")
            if event.content.parts:
                response_text = event.content.parts[0].text
                print(f"   {response_text}")
                if event.is_final_response():
                    final_response = response_text
    session = await session_service.get_session(
        app_name=APP_NAME,
        user_id=USER_ID, 
        session_id=session_id,
    )
    print(f"{'='*60}")
    print(f"📊 Workflow Complete - Session State ({session_id}):")
    print(f"{'='*60}")
    for key, value in session.state.items():
        print(f"   {key}: {value}")
    print(f"🎯 Final Response: {final_response}")
    return final_response

async def main_notebook():
    test_requests = [
        "I need approval for $750 for team lunch and celebrations",
        "Please approve $3,000 for a conference ticket and travel expenses",
        "I need $12,000 approved for critical software licenses renewal"
    ]
    for i, request in enumerate(test_requests, 1):
        current_session_id = f"approval_session_notebook_{456 + i -1}"
        # Create the session before running the workflow
        await session_service.create_session(
            app_name=APP_NAME,
            user_id=USER_ID,
            session_id=current_session_id
        )
        print(f"Created session: {current_session_id}")
        await run_approval_workflow_notebook(request, current_session_id)
try:
    asyncio.run(main_notebook())
    agentops.end_trace(end_state="Success")
except Exception as e:
    print(f"Error: {e}")
    agentops.end_trace(end_state="Error")