import traceback
from agents import RunContextWrapper, Runner, GuardrailFunctionOutput, input_guardrail
from guardrails.input_guardrails import guardrail_agent

@input_guardrail
async def guardrail_input_function(ctx: RunContextWrapper, agent, user_input: str):
    try:
        result = await Runner.run(
            guardrail_agent, 
            input=user_input, 
            context=ctx.context
        )
        
        # Check if result has the expected structure
        if not result or not hasattr(result, 'final_output'):
            print(f"Warning: Guardrail agent returned unexpected result: {result}")
            # Allow the query to proceed if guardrail fails
            return GuardrailFunctionOutput(
                output_info=None,
                tripwire_triggered=False
            )
        
        final_output = result.final_output
        
        # Check if final_output has the expected attribute
        if not hasattr(final_output, 'is_query_about_launchlabs'):
            print(f"Warning: Guardrail output missing is_query_about_launchlabs attribute: {final_output}")
            return GuardrailFunctionOutput(
                output_info=final_output,
                tripwire_triggered=False
            )
        
        return GuardrailFunctionOutput(
            output_info=final_output,
            tripwire_triggered=not final_output.is_query_about_launchlabs
        )
    except Exception as e:
        print(f"Error in guardrail_input_function: {e}")
        print(traceback.format_exc())
        # Allow the query to proceed if guardrail fails
        return GuardrailFunctionOutput(
            output_info=None,
            tripwire_triggered=False
        )