import { runAgent, RunStateStore } from '@timestep-ai/timestep';
import {
  Agent,
  OpenAIConversationsSession,
  tool,
  InputGuardrail,
  OutputGuardrail,
} from '@openai/agents';
import type { AgentInputItem } from '@openai/agents-core';

// Define a tool with approval requirement
const getWeather = tool({
  name: 'get_weather',
  description: 'Get weather information for a city.',
  parameters: {
    type: 'object',
    properties: {
      city: { type: 'string' }
    },
    required: ['city'],
    additionalProperties: false
  } as any,
  needsApproval: async (_ctx: any, args: any) => {
    // Require approval for sensitive cities
    return ['berkeley', 'san francisco'].includes(args.city?.toLowerCase() || '');
  },
  execute: async (args: any): Promise<string> => {
    return `The weather in ${args.city} is sunny and 72°F`;
  }
});

// Define guardrails
const contentFilter: InputGuardrail = {
  name: 'Content Filter',
  runInParallel: true,
  execute: async ({ input }) => {
    const inputText = typeof input === 'string' ? input : JSON.stringify(input);
    const blocked = ['spam', 'scam'].some(word => 
      inputText.toLowerCase().includes(word)
    );
    return {
      outputInfo: { blocked },
      tripwireTriggered: blocked,
    };
  },
};

const outputSafety: OutputGuardrail<any> = {
  name: 'Output Safety',
  execute: async ({ agentOutput }) => {
    const outputText = typeof agentOutput === 'string' 
      ? agentOutput 
      : JSON.stringify(agentOutput);
    const unsafe = outputText.toLowerCase().includes('password') || 
                   outputText.toLowerCase().includes('api key');
    return {
      outputInfo: { unsafe },
      tripwireTriggered: unsafe,
    };
  },
};

// Create specialized weather agent
const weatherAgent = new Agent({
  instructions: 'You are a weather assistant. Use get_weather for all weather queries.',
  model: 'gpt-4.1',
  modelSettings: { temperature: 0 },
  name: 'Weather Assistant',
  tools: [getWeather],
});

// Create main assistant with handoffs and guardrails
const assistant = new Agent({
  handoffs: [weatherAgent],
  instructions: 'You are a helpful personal assistant.',
  model: 'gpt-4.1',
  modelSettings: { temperature: 0 },
  name: 'Personal Assistant',
  inputGuardrails: [contentFilter],
  outputGuardrails: [outputSafety],
});

// Initialize session and state store
const session = new OpenAIConversationsSession();
const sessionId = await session.getSessionId();
const stateStore = new RunStateStore({ agent: assistant, sessionId });

// Run multiple conversation turns
const conversations: AgentInputItem[][] = [
  [{ type: 'message', role: 'user', content: [{ type: 'input_text', text: "What's 2+2?" }] }],
  [{ type: 'message', role: 'user', content: [{ type: 'input_text', text: "What's the weather in Oakland?" }] }],
  [{ type: 'message', role: 'user', content: [{ type: 'input_text', text: "What's the weather in Berkeley?" }] }],
];

for (const inputItems of conversations) {
  let result = await runAgent(assistant, inputItems, session, false);
  
  // Handle interruptions (e.g., tool approval needed)
  if (result.interruptions?.length) {
    await stateStore.save(result.state);
    
    // Load, approve, and resume
    const loadedState = await stateStore.load();
    for (const interruption of loadedState.getInterruptions()) {
      loadedState.approve(interruption);
    }
    
    result = await runAgent(assistant, loadedState, session, false);
  }
}

const saved = await stateStore.load();
for (const interruption of saved.getInterruptions()) saved.approve(interruption);
await runAgent(agent, saved, session, false);