// src/lib/agent.ts
export interface Step {
  number: number;
  description: string;
  reasoning?: string;
  output?: string;
}

export interface AgentResult {
  goal: string;
  steps: Step[];
  output: string;
}

export type StepCallback = (step: Step) => Promise<void> | void;

export class Agent {
  private goal: string;
  private maxSteps: number;
  private onStepComplete?: StepCallback;
  private steps: Step[] = [];

  constructor(options: {
    goal: string;
    maxSteps?: number;
    onStepComplete?: StepCallback;
  }) {
    this.goal = options.goal;
    this.maxSteps = options.maxSteps || 5;
    this.onStepComplete = options.onStepComplete;
  }

  async execute(): Promise<AgentResult> {
    // Step 1: Task analysis
    const taskAnalysis = await this.callOllama(
      `Analyze this task: "${this.goal}". Break it down into ${this.maxSteps} clear steps that would lead to a high-quality result. Return a JSON array of step descriptions only, no additional text.`
    );
    
    let steps: string[] = [];
    try {
      const parsed = JSON.parse(this.extractJSON(taskAnalysis));
      steps = Array.isArray(parsed) ? parsed : [];
    } catch (e) {
      // If parsing fails, try to extract steps using regex
      const stepRegex = /\d+\.\s*(.*?)(?=\d+\.|$)/gs;
      const matches = [...taskAnalysis.matchAll(stepRegex)];
      steps = matches.map(match => match[1].trim());
    }
    
    // Ensure we have steps
    if (steps.length === 0) {
      steps = ["Analyze the problem", "Generate solution", "Refine the output"];
    }
    
    // Execute each step
    for (let i = 0; i < Math.min(steps.length, this.maxSteps); i++) {
      const stepNumber = i + 1;
      const stepDescription = steps[i];
      
      // Generate reasoning for this step
      const reasoning = await this.callOllama(
        `For the task: "${this.goal}", I am on step ${stepNumber}: "${stepDescription}". Explain your reasoning for how you'll approach this step. Keep it clear and concise.`
      );
      
      // Execute the step
      const stepPrompt = `
Task: "${this.goal}"
Step ${stepNumber}/${Math.min(steps.length, this.maxSteps)}: ${stepDescription}
Previous steps: ${this.steps.map(s => `Step ${s.number}: ${s.description} -> ${s.output?.substring(0, 100)}...`).join('\n')}

Execute this step and provide the output. Be thorough but focused on just this step.
`;
      
      const stepOutput = await this.callOllama(stepPrompt);
      
      // Record the step
      const step: Step = {
        number: stepNumber,
        description: stepDescription,
        reasoning,
        output: stepOutput
      };
      
      this.steps.push(step);
      
      // Notify via callback if provided
      if (this.onStepComplete) {
        await this.onStepComplete(step);
      }
    }
    
    // Generate final comprehensive output
    const finalPrompt = `
You've been working on: "${this.goal}"

You've completed the following steps:
${this.steps.map(s => `Step ${s.number}: ${s.description}`).join('\n')}

Now, compile all of your work into a comprehensive final output that achieves the original goal. 
Format your response using Markdown for readability. Include headings, bullet points, and other formatting as appropriate.
Ensure your response is complete, well-structured, and directly addresses the original goal.
`;
    
    const finalOutput = await this.callOllama(finalPrompt);
    
    return {
      goal: this.goal,
      steps: this.steps,
      output: finalOutput
    };
  }

  private async callOllama(prompt: string): Promise<string> {
    try {
      const response = await fetch('http://localhost:11434/api/generate', {
        method: 'POST',
        headers: {
          'Content-Type': 'application/json',
        },
        body: JSON.stringify({
          model: 'mistral',
          prompt: prompt,
          stream: false,
        }),
      });

      if (!response.ok) {
        throw new Error(`Ollama API error: ${response.statusText}`);
      }

      const data = await response.json();
      return data.response;
    } catch (error) {
      console.error('Error calling Ollama:', error);
      return `Error: ${error instanceof Error ? error.message : 'Unknown error'}`;
    }
  }

  private extractJSON(text: string): string {
    // Try to extract JSON from the text
    const jsonRegex = /(\[.*\]|\{.*\})/s;
    const match = text.match(jsonRegex);
    return match ? match[0] : '[]';
  }
}

import { NextResponse } from 'next/server';
import { Agent, Step } from '@/lib/agent';

export const runtime = 'nodejs';

export async function POST(request: Request) {
  // Initialize the response encoder for streaming
  const encoder = new TextEncoder();
  const stream = new TransformStream();
  const writer = stream.writable.getWriter();

  // Function to send updates to the client
  const sendUpdate = async (data: any) => {
    await writer.write(encoder.encode(JSON.stringify(data) + '\n'));
  };

  // Process the request in the background while streaming updates
  const processRequest = async () => {
    try {
      // Parse the request body
      const { goal } = await request.json();
      
      if (!goal || typeof goal !== 'string') {
        await sendUpdate({
          type: 'log',
          message: '❌ Error: Please provide a valid goal'
        });
        writer.close();
        return;
      }

      // Initialize the custom agent
      const agent = new Agent({
        goal: goal,
        maxSteps: 5,
        onStepComplete: async (step: Step) => {
          await sendUpdate({
            type: 'log',
            message: `📝 Step ${step.number}: ${step.description}`
          });
          
          if (step.reasoning) {
            await sendUpdate({
              type: 'log',
              message: `🤔 Reasoning: ${step.reasoning}`
            });
          }
        }
      });
      
      // Log initialization
      await sendUpdate({
        type: 'log',
        message: `🧠 Analyzing task: "${goal}"`
      });
      
      // Execute the agent
      const result = await agent.execute();
      
      // Send the final result
      await sendUpdate({
        type: 'result',
        content: result.output
      });
      
      // Close the stream
      writer.close();
    } catch (error: any) {
      console.error('Agent execution error:', error);
      await sendUpdate({
        type: 'log',
        message: `❌ Error: ${error.message || 'Unknown error occurred'}`
      });
      writer.close();
    }
  };
  
  // Start processing in the background
  processRequest();
  
  // Return the stream response immediately
  return new NextResponse(stream.readable, {
    headers: {
      'Content-Type': 'application/json',
      'Transfer-Encoding': 'chunked',
    },
  });
}

// src/lib/tools.ts
export interface Tool {
  name: string;
  description: string;
  execute: (input: string) => Promise<string>;
}

// Sample search tool
export const searchTool: Tool = {
  name: 'search',
  description: 'Search the web for information',
  async execute(query: string): Promise<string> {
    // This is a mock implementation - in a real app, you'd integrate with a search API
    return `Simulated search results for: ${query}\n\n1. First relevant result\n2. Second relevant result\n3. Third relevant result`;
  }
};

// Then enhance your Agent class to use tools
// In src/lib/agent.ts, modify the execute method:

async execute(): Promise<AgentResult> {
  // ... existing code
  
  // Add tool usage where appropriate
  if (this.tools && this.tools.length > 0) {
    const toolDescriptions = this.tools.map(t => `${t.name}: ${t.description}`).join('\n');
    
    // Let the agent decide whether to use a tool
    const toolUsage = await this.callOllama(`
      For the task: "${this.goal}", I'm on step ${currentStep.number}: "${currentStep.description}".
      I have these tools available:
      ${toolDescriptions}
      
      Should I use a tool for this step? If yes, specify which tool and the exact input to provide to the tool.
      Return your answer in JSON format: { "useTool": boolean, "toolName": string, "toolInput": string }
    `);
    
    try {
      const toolDecision = JSON.parse(this.extractJSON(toolUsage));
      if (toolDecision.useTool) {
        const tool = this.tools.find(t => t.name === toolDecision.toolName);
        if (tool) {
          const toolResult = await tool.execute(toolDecision.toolInput);
          // Use the tool result in further processing
          currentStep.output = `Used ${tool.name} with input: ${toolDecision.toolInput}\n\nResult: ${toolResult}`;
        }
      }
    } catch (e) {
      // If JSON parsing fails, continue without tool usage
    }
  }
  
  // ... rest of execute method
}

// Using Prisma with SQLite for simplicity
// 1. Install Prisma: npm install prisma @prisma/client
// 2. Initialize Prisma: npx prisma init --datasource-provider sqlite

// 3. Create schema.prisma:
// prisma/schema.prisma
/*
generator client {
  provider = "prisma-client-js"
}

datasource db {
  provider = "sqlite"
  url      = "file:./dev.db"
}

model Session {
  id        String   @id @default(uuid())
  goal      String
  output    String   @default("")
  logs      String   @default("[]") // JSON string of logs
  createdAt DateTime @default(now())
}
*/

// 4. Run migration: npx prisma migrate dev --name init
// 5. Generate client: npx prisma generate

// 6. Create a database service
// src/lib/db.ts
import { PrismaClient } from '@prisma/client';

const prisma = new PrismaClient();

export async function saveSession(goal: string, output: string, logs: string[]) {
  return await prisma.session.create({
    data: {
      goal,
      output,
      logs: JSON.stringify(logs),
    },
  });
}

export async function getSessions() {
  return await prisma.session.findMany({
    orderBy: {
      createdAt: 'desc',
    },
  });
}

export async function getSession(id: string) {
  return await prisma.session.findUnique({
    where: { id },
  });
}

// src/lib/workflow.ts
import { Agent, AgentResult } from './agent';

export async function runResearchAndSynthesisWorkflow(topic: string, updateCallback: (message: string) => Promise<void>) {
  await updateCallback(`Starting research workflow on: ${topic}`);
  
  // Research agent gathers information
  const researchAgent = new Agent({
    goal: `Research key facts about: ${topic}`,
    maxSteps: 3,
    onStepComplete: async (step) => {
      await updateCallback(`Research step ${step.number}: ${step.description}`);
    }
  });
  
  await updateCallback("Starting research phase...");
  const researchResult = await researchAgent.execute();
  
  // Analysis agent evaluates the research
  const analysisAgent = new Agent({
    goal: `Analyze these research findings and identify key insights: ${researchResult.output}`,
    maxSteps: 2,
    onStepComplete: async (step) => {
      await updateCallback(`Analysis step ${step.number}: ${step.description}`);
    }
  });
  
  await updateCallback("Starting analysis phase...");
  const analysisResult = await analysisAgent.execute();
  
  // Synthesis agent creates final output
  const synthesisAgent = new Agent({
    goal: `Create a comprehensive report on ${topic} using this research and analysis:
    Research: ${researchResult.output}
    Analysis: ${analysisResult.output}`,
    maxSteps: 3,
    onStepComplete: async (step) => {
      await updateCallback(`Synthesis step ${step.number}: ${step.description}`);
    }
  });
  
  await updateCallback("Starting synthesis phase...");
  const finalResult = await synthesisAgent.execute();
  
  await updateCallback("Workflow complete!");
  
  return {
    topic,
    research: researchResult.output,
    analysis: analysisResult.output,
    synthesis: finalResult.output
  };
}

// src/lib/modelConfig.ts
export interface ModelConfig {
  modelName: string;
  temperature: number;
  maxTokens: number;
}

export const availableModels = [
  { name: 'mistral', label: 'Mistral (Balanced)' },
  { name: 'llama3', label: 'Llama 3 (Creative)' },
  { name: 'codellama', label: 'CodeLlama (Technical)' },
];

// Then update your Agent class to use these configurations:
private async callOllama(prompt: string): Promise<string> {
  try {
    const response = await fetch('http://localhost:11434/api/generate', {
      method: 'POST',
      headers: {
        'Content-Type': 'application/json',
      },
      body: JSON.stringify({
        model: this.modelConfig.modelName || 'mistral',
        prompt: prompt,
        temperature: this.modelConfig.temperature || 0.7,
        max_tokens: this.modelConfig.maxTokens || 2048,
        stream: false,
      }),
    });

    if (!response.ok) {
      throw new Error(`Ollama API error: ${response.statusText}`);
    }

    const data = await response.json();
    return data.response;
  } catch (error) {
    console.error('Error calling Ollama:', error);
    return `Error: ${error instanceof Error ? error.message : 'Unknown error'}`;
  }
}