// agent.ts
import Anthropic from "@anthropic-ai/sdk";
import { z } from "zod";

const client = new Anthropic();

// Define the tools this agent can use
const tools: Anthropic.Tool[] = [
  {
    name: "search_documents",
    description: "Search internal knowledge base for relevant documents",
    input_schema: {
      type: "object",
      properties: {
        query: { type: "string", description: "Search query" },
        limit: { type: "number", description: "Max results to return", default: 5 },
      },
      required: ["query"],
    },
  },
  {
    name: "create_summary",
    description: "Persist a structured summary to the database",
    input_schema: {
      type: "object",
      properties: {
        title: { type: "string" },
        key_points: { type: "array", items: { type: "string" } },
        confidence: { type: "number", minimum: 0, maximum: 1 },
      },
      required: ["title", "key_points", "confidence"],
    },
  },
];

// Structured output schema — enforced after the agent finishes
const AgentResultSchema = z.object({
  summary_id: z.string(),
  topics_covered: z.array(z.string()),
  gaps_identified: z.array(z.string()),
});

type AgentResult = z.infer<typeof AgentResultSchema>;

async function runAgent(userQuery: string): Promise<AgentResult> {
  const messages: Anthropic.MessageParam[] = [
    { role: "user", content: userQuery },
  ];

  // Agentic loop — continues until the model stops calling tools
  while (true) {
    const response = await client.messages.create({
      model: "claude-opus-4-6",
      max_tokens: 4096,
      tools,
      messages,
      system: `You are a research assistant. Use the available tools to answer 
               the user's query thoroughly. Always create a summary when you 
               have enough information. Be explicit about gaps in your knowledge.`,
    });

    // Append assistant turn to message history
    messages.push({ role: "assistant", content: response.content });

    // If no more tool calls, we're done
    if (response.stop_reason === "end_turn") {
      break;
    }

    // Process tool calls
    const toolResults: Anthropic.ToolResultBlockParam[] = [];

    for (const block of response.content) {
      if (block.type !== "tool_use") continue;

      const result = await executeTool(block.name, block.input as Record<string, unknown>);
      toolResults.push({
        type: "tool_result",
        tool_use_id: block.id,
        content: JSON.stringify(result),
      });
    }

    // Feed results back into the loop
    messages.push({ role: "user", content: toolResults });
  }

  // Extract and validate the final structured result
  const lastMessage = messages[messages.length - 1];
  const rawOutput = extractJSON(lastMessage);
  return AgentResultSchema.parse(rawOutput);
}

// Stub — in production this dispatches to your actual tool implementations
async function executeTool(name: string, input: Record<string, unknown>): Promise<unknown> {
  console.log(`[tool] ${name}`, input);
  // Replace with real implementations
  return { status: "ok", data: [] };
}

function extractJSON(message: Anthropic.MessageParam): unknown {
  const content = Array.isArray(message.content)
    ? message.content.find((b) => b.type === "text")?.text ?? ""
    : message.content;
  const match = content.match(/```json\n([\s\S]+?)\n```/);
  if (!match) throw new Error("No JSON block found in agent response");
  return JSON.parse(match[1]);
}

// orchestrator.ts
import { Worker } from "worker_threads";
import path from "path";

interface DocumentJob {
  id: string;
  content: string;
  metadata: Record<string, unknown>;
}

interface DocumentResult {
  id: string;
  entities: string[];
  sentiment: "positive" | "negative" | "neutral";
  summary: string;
}

const MAX_CONCURRENT_AGENTS = 5; // Respect API rate limits

async function processDocumentBatch(
  documents: DocumentJob[]
): Promise<DocumentResult[]> {
  const results: DocumentResult[] = [];
  const queue = [...documents];
  const inFlight: Promise<DocumentResult>[] = [];

  while (queue.length > 0 || inFlight.length > 0) {
    // Fill up to the concurrency limit
    while (queue.length > 0 && inFlight.length < MAX_CONCURRENT_AGENTS) {
      const job = queue.shift()!;
      const promise = processDocument(job).then((result) => {
        // Remove from in-flight when done
        inFlight.splice(inFlight.indexOf(promise), 1);
        return result;
      });
      inFlight.push(promise);
    }

    // Wait for at least one to finish before looping
    if (inFlight.length > 0) {
      const result = await Promise.race(inFlight);
      results.push(result);
    }
  }

  return results;
}

async function processDocument(job: DocumentJob): Promise<DocumentResult> {
  // Each document gets its own isolated agent call
  // In production this might be a separate process, Lambda, or worker thread
  const response = await fetch("http://localhost:8000/process", {
    method: "POST",
    headers: { "Content-Type": "application/json" },
    body: JSON.stringify(job),
  });

  if (!response.ok) {
    throw new Error(`Agent failed for document ${job.id}: ${response.statusText}`);
  }

  return response.json() as Promise<DocumentResult>;
}

# worker.py
from fastapi import FastAPI, HTTPException
from anthropic import Anthropic
from pydantic import BaseModel
from typing import Literal
import json

app = FastAPI()
client = Anthropic()


class DocumentJob(BaseModel):
    id: str
    content: str
    metadata: dict


class DocumentResult(BaseModel):
    id: str
    entities: list[str]
    sentiment: Literal["positive", "negative", "neutral"]
    summary: str


@app.post("/process", response_model=DocumentResult)
async def process_document(job: DocumentJob) -> DocumentResult:
    """Each HTTP request is an isolated agent. No shared state."""
    
    response = client.messages.create(
        model="claude-opus-4-6",
        max_tokens=1024,
        system="""You are a document analysis agent. Analyze the given document and 
                  return a JSON object with exactly these fields:
                  - entities: list of named entities (people, orgs, places)
                  - sentiment: one of "positive", "negative", "neutral"  
                  - summary: one sentence summary
                  
                  Respond ONLY with valid JSON. No prose.""",
        messages=[
            {
                "role": "user",
                "content": f"Document ID: {job.id}\n\n{job.content}"
            }
        ],
    )

    raw = response.content[0].text.strip()
    
    try:
        data = json.loads(raw)
        return DocumentResult(id=job.id, **data)
    except (json.JSONDecodeError, KeyError) as e:
        raise HTTPException(
            status_code=500,
            detail=f"Agent returned invalid output for document {job.id}: {e}"
        )

# Start the document processing worker
uvicorn worker:app --host 0.0.0.0 --port 8000 --workers 4

# Or with Docker for isolation (recommended in production)
docker run -d \
  --name doc-processor \
  -p 8000:8000 \
  -e ANTHROPIC_API_KEY=$ANTHROPIC_API_KEY \
  --restart unless-stopped \
  everscape/doc-processor:latest