// app/api/agent/route.ts
export async function POST(req: Request) {
  const { goal } = await req.json();
  const encoder = new TextEncoder();
  
  const stream = new ReadableStream({
    async start(controller) {
      // Callback that sends tokens as they're generated
      const sendToken = (token: string) => {
        controller.enqueue(encoder.encode(token));
      };
      
      try {
        // Main agent execution with streaming callback
        await runAgent({ goal }, sendToken);
      } catch (error) {
        sendToken(`\nError: ${error.message}`);
      } finally {
        controller.close();
      }
    },
  });

  return new Response(stream, {
    headers: {
      'Content-Type': 'text/event-stream',
      'Cache-Control': 'no-cache',
      'Connection': 'keep-alive',
    },
  });
}

// Example of model selection logic
function selectModelForTask(task) {
  if (task.type === 'code_generation') {
    return 'codellama';
  } else if (task.complexity === 'high') {
    return 'llama2:70b';
  } else if (task.priority === 'speed') {
    return 'mistral';
  } else {
    return 'llama2:13b'; // Default
  }
}

import { v4 as uuid } from 'uuid';
import { decomposeTask } from './decompose';
import { executeStep } from './execute';

export async function runAgent(
  { goal }: { goal: string },
  streamCallback: (text: string) => void
) {
  const sessionId = uuid();
  
  try {
    // Start by streaming a confirmation
    streamCallback(`📋 Working on: "${goal}"\n\n`);
    
    // Step 1: Decompose the task
    streamCallback(`🔍 Analyzing task and creating plan...\n`);
    const steps = await decomposeTask(goal);
    
    streamCallback(`✅ Created plan with ${steps.length} steps:\n\n`);
    steps.forEach((step, i) => {
      streamCallback(`${i+1}. ${step.description}\n`);
    });
    streamCallback(`\n`);
    
    // Step 2: Execute each step with proper dependency handling
    const results = {};
    const completed = new Set();
    
    // Keep going until all steps are complete
    while (completed.size < steps.length) {
      // Find steps where all dependencies are satisfied
      const availableSteps = steps.filter(step => {
        if (completed.has(step.id)) return false;
        
        // Check if all dependencies are completed
        return step.depends_on.every(depId => completed.has(depId));
      });
      
      if (availableSteps.length === 0) {
        throw new Error("Deadlock detected - unable to proceed with execution plan");
      }
      
      // Execute available steps in parallel
      streamCallback(`🔄 Executing ${availableSteps.length} steps in parallel...\n\n`);
      
      await Promise.all(availableSteps.map(async (step) => {
        streamCallback(`⏳ Working on: ${step.description}...\n`);
        
        try {
          // Get context from dependencies
          const context = {};
          step.depends_on.forEach(depId => {
            context[depId] = results[depId];
          });
          
          // Execute the step
          const result = await executeStep(step, context, goal);
          results[step.id] = result;
          completed.add(step.id);
          
          streamCallback(`✅ Completed: ${step.description}\n`);
          streamCallback(`   Result: ${result.substring(0, 100)}${result.length > 100 ? '...' : ''}\n\n`);
        } catch (error) {
          streamCallback(`❌ Error in step "${step.description}": ${error.message}\n\n`);
          throw error;
        }
      }));
    }
    
    // Step 3: Generate final summary
    streamCallback(`🎉 All steps completed! Generating final report...\n\n`);
    
    const summaryPrompt = `
      You've completed all steps to: ${goal}
      
      Here are the results from each step:
      ${Object.entries(results).map(([id, result]) => {
        const step = steps.find(s => s.id === id);
        return `STEP: ${step.description}\nRESULT: ${result}\n\n`;
      }).join('')}
      
      Please provide a comprehensive final report that synthesizes all this information.
    `;
    
    const summary = await generateWithOllama('llama2', summaryPrompt);
    streamCallback(`📊 FINAL REPORT:\n\n${summary}\n\n`);
    
    return { sessionId, success: true };
  } catch (error) {
    streamCallback(`\n❌ Agent execution failed: ${error.message}\n`);
    return { sessionId, success: false, error: error.message };
  }
}

// Helper function for Ollama API calls
async function generateWithOllama(model, prompt) {
  const response = await fetch('http://localhost:11434/api/generate', {
    method: 'POST',
    headers: { 'Content-Type': 'application/json' },
    body: JSON.stringify({ model, prompt, stream: false }),
  });
  
  if (!response.ok) {
    const error = await response.text();
    throw new Error(`Ollama error: ${error}`);
  }
  
  const data = await response.json();
  return data.response;
}

import { generateWithOllama } from '../ollama';
import { v4 as uuid } from 'uuid';

export async function decomposeTask(goal: string) {
  const decompositionPrompt = `
    You are an expert project manager who breaks down complex tasks into well-organized steps.
    
    RULES:
    1. Create 3-7 logical steps to accomplish the goal
    2. Each step should be self-contained and focused
    3. Identify dependencies between steps (which steps must complete before others)
    4. Steps with no dependencies should be executable in parallel
    
    GOAL: ${goal}
    
    Respond in valid JSON format:
    {
      "steps": [
        {
          "id": "unique_id_1",
          "description": "Clear description of the step",
          "depends_on": [] // Array of step IDs this step depends on
        },
        ...
      ]
    }
  `;
  
  try {
    const result = await generateWithOllama('llama2', decompositionPrompt);
    let parsed;
    
    try {
      // Find the JSON part in the response (models sometimes add explanations)
      const jsonMatch = result.match(/\{[\s\S]*\}/);
      parsed = JSON.parse(jsonMatch ? jsonMatch[0] : result);
    } catch (e) {
      throw new Error(`Failed to parse task decomposition: ${e.message}`);
    }
    
    if (!parsed.steps || !Array.isArray(parsed.steps)) {
      throw new Error('Invalid decomposition format: missing steps array');
    }
    
    // Assign IDs if they're missing
    parsed.steps = parsed.steps.map(step => ({
      ...step,
      id: step.id || `step_${uuid().substring(0, 8)}`,
      depends_on: step.depends_on || []
    }));
    
    return parsed.steps;
  } catch (error) {
    console.error('Decomposition error:', error);
    throw new Error(`Failed to decompose task: ${error.message}`);
  }
}

import { generateWithOllama } from '../ollama';

export async function executeStep(step, context, originalGoal) {
  // Build context from dependent steps
  const dependencyContext = step.depends_on.map(depId => {
    const result = context[depId];
    return `Previous Step Result (${depId}): ${result}`;
  }).join('\n\n');
  
  const executionPrompt = `
    You are an AI assistant working on a multi-step task.
    
    ORIGINAL GOAL: ${originalGoal}
    
    CURRENT STEP: ${step.description}
    
    ${dependencyContext ? `CONTEXT FROM PREVIOUS STEPS:\n${dependencyContext}\n\n` : ''}
    
    Your task is to complete ONLY this specific step thoroughly and accurately.
    Provide a complete solution for this step only.
  `;
  
  try {
    return await generateWithOllama('llama2', executionPrompt);
  } catch (error) {
    throw new Error(`Step execution failed: ${error.message}`);
  }
}

import { NextRequest } from 'next/server';
import { runAgent } from '@/lib/agent';

export async function POST(req: NextRequest) {
  try {
    const { goal } = await req.json();
    
    if (!goal || typeof goal !== 'string') {
      return new Response(
        JSON.stringify({ error: 'Missing or invalid goal' }),
        { status: 400, headers: { 'Content-Type': 'application/json' } }
      );
    }
    
    // Create encoder for streaming
    const encoder = new TextEncoder();
    
    // Create a stream for the response
    const stream = new ReadableStream({
      async start(controller) {
        const sendToken = (token: string) => {
          controller.enqueue(encoder.encode(token));
        };
        
        try {
          await runAgent({ goal }, sendToken);
        } catch (error) {
          sendToken(`\nError: ${error.message}`);
          console.error('Agent execution error:', error);
        } finally {
          controller.close();
        }
      },
    });
    
    // Return the stream as a streamed response
    return new Response(stream, {
      headers: {
        'Content-Type': 'text/event-stream',
        'Cache-Control': 'no-cache',
        'Connection': 'keep-alive',
      },
    });
  } catch (error) {
    console.error('API route error:', error);
    return new Response(
      JSON.stringify({ error: 'Internal server error' }),
      { status: 500, headers: { 'Content-Type': 'application/json' } }
    );
  }
}

export const PROMPTS = {
  // For travel planning domain
  travel: {
    decompose: `
      You are a professional travel planner. Break down the travel planning process into logical steps.
      
      Consider:
      - Transportation research
      - Accommodation booking
      - Activity planning
      - Budget management
      - Visa/document requirements
      
      GOAL: {goal}
      
      Create a comprehensive plan in JSON format:
      {
        "steps": [...]
      }
    `,
    
    // Other travel-specific prompts
  },
  
  // For research paper domain
  research: {
    decompose: `
      You are a research assistant helping to plan an academic paper. Break down the research process.
      
      Consider:
      - Literature review
      - Hypothesis formation
      - Methodology planning
      - Data collection strategy
      - Analysis approach
      
      GOAL: {goal}
      
      Create a comprehensive plan in JSON format:
      {
        "steps": [...]
      }
    `,
    
    // Other research-specific prompts
  }
};

export function getPrompt(domain, type, variables) {
  const template = PROMPTS[domain]?.[type] || PROMPTS.general[type];
  
  return Object.entries(variables).reduce((prompt, [key, value]) => {
    return prompt.replace(new RegExp(`{${key}}`, 'g'), value);
  }, template);
}

export const tools = {
  // Get weather information
  weather: async (location) => {
    try {
      // Using a local-only weather API
      const res = await fetch(`http://localhost:3001/api/weather?location=${encodeURIComponent(location)}`);
      if (!res.ok) throw new Error('Weather service unavailable');
      return await res.json();
    } catch (error) {
      return { error: error.message };
    }
  },
  
  // Calculate distance between locations
  distance: (origin, destination) => {
    // Simple demo implementation (would use local libraries in real app)
    return { 
      distance: "1,234 km",
      duration: "Approximately 13 hours driving"
    };
  },
  
  // Current date/time (no external API needed)
  datetime: () => {
    const now = new Date();
    return {
      date: now.toLocaleDateString(),
      time: now.toLocaleTimeString(),
      timestamp: now.getTime()
    };
  }
};

// Tool usage function for the agent
export async function useTool(toolName, ...args) {
  if (!tools[toolName]) {
    return { error: `Tool '${toolName}' not found` };
  }
  
  try {
    return await tools[toolName](...args);
  } catch (error) {
    return { error: error.message };
  }
}

export async function withErrorHandling(fn, options = {}) {
  const { 
    maxRetries = 3, 
    retryDelay = 1000,
    fallback = null
  } = options;
  
  let lastError;
  
  for (let attempt = 1; attempt <= maxRetries; attempt++) {
    try {
      return await fn();
    } catch (error) {
      lastError = error;
      
      // Log the error with attempt count
      console.error(`Attempt ${attempt}/${maxRetries} failed:`, error);
      
      // Check if we should retry based on error type
      if (error.message.includes('context length exceeded')) {
        console.log('Context length error - trying with truncated input');
        // We could modify fn here to use less context
      } else if (error.message.includes('connection failed')) {
        console.log(`Connection error - retrying in ${retryDelay}ms`);
        await new Promise(r => setTimeout(r, retryDelay));
      } else if (attempt >= maxRetries) {
        // If this is the last attempt, don't retry
        break;
      }
    }
  }
  
  // All attempts failed
  if (fallback) {
    console.log('All attempts failed, using fallback');
    return fallback();
  }
  
  throw lastError;
}