import { NextRequest } from "next/server"
import Anthropic from "@anthropic-ai/sdk"
import { getTraceData } from "@/app/observability/lib/get-trace-data"
import {
  indexTraceData,
  buildSummaryPrompt,
  anthropicToolDefinitions,
  resolveToolCall,
  type IndexedTraceData,
} from "@/app/observability/lib/build-chat-context"

interface ChatMessage {
  role: "user" | "assistant"
  content: string
}

const MAX_CLIENT_CONTENT = 4000
const KEEPALIVE_INTERVAL_MS = 15_000
const ROUND_TIMEOUT_MS = 3 * 60_000 // 3 minutes per API round

const TOOL_DISPLAY_NAMES: Record<string, string> = {
  get_original_code: "Reading original code",
  get_candidate_code: "Reading candidate code",
  get_test_code: "Reading test code",
  get_ranking_details: "Reading ranking details",
  get_errors: "Checking errors",
  get_llm_call_detail: "Inspecting LLM call",
  read_file: "Reading file",
  search_code: "Searching codebase",
  list_directory: "Listing directory",
}

function getClient(): Anthropic {
  const baseURL = process.env.ANTHROPIC_FOUNDRY_BASE_URL
  const apiKey = process.env.AZURE_OPENAI_API_KEY
  if (!baseURL || !apiKey) {
    throw new Error("ANTHROPIC_FOUNDRY_BASE_URL and AZURE_OPENAI_API_KEY must be configured")
  }
  return new Anthropic({ baseURL, apiKey })
}

function truncateForClient(content: string): string {
  if (content.length <= MAX_CLIENT_CONTENT) return content
  return content.slice(0, MAX_CLIENT_CONTENT) + "\n\n... [truncated]"
}

function summarizeToolResult(toolName: string, result: string): string {
  const lines = result.split("\n").filter(Boolean)
  switch (toolName) {
    case "search_code": {
      if (result === "No matches found.") return "No matches found"
      return `Found ${lines.length} matching lines`
    }
    case "read_file": {
      const header = lines[0] ?? ""
      return header.startsWith("File:") ? header : `Read ${lines.length} lines`
    }
    case "get_errors": {
      if (result === "No errors in this trace.") return "No errors"
      const count = lines.filter(l => l.startsWith("[")).length
      return `Found ${count} errors`
    }
    case "get_llm_call_detail":
      return "Loaded full LLM call details"
    case "list_directory": {
      const count = lines.length - 1 // subtract header line
      return `Listed ${count} entries`
    }
    default:
      return "Done"
  }
}

async function processToolCalls(
  content: Anthropic.ContentBlock[],
  indexed: IndexedTraceData,
  enqueue: (data: string) => void,
): Promise<Anthropic.ToolResultBlockParam[]> {
  const toolUseBlocks = content.filter(
    (block): block is Anthropic.ToolUseBlock => block.type === "tool_use",
  )
  if (toolUseBlocks.length === 0) return []

  for (const block of toolUseBlocks) {
    enqueue(
      `data: ${JSON.stringify({
        type: "tool_start",
        tool: block.name,
        displayName: TOOL_DISPLAY_NAMES[block.name] ?? block.name,
      })}\n\n`,
    )
  }

  const results = await Promise.all(
    toolUseBlocks.map(async block => {
      const result = await resolveToolCall(
        block.name,
        (block.input as Record<string, unknown>) ?? {},
        indexed,
      )

      enqueue(
        `data: ${JSON.stringify({
          type: "tool_result",
          tool: block.name,
          displayName: TOOL_DISPLAY_NAMES[block.name] ?? block.name,
          summary: summarizeToolResult(block.name, result),
          content: truncateForClient(result),
        })}\n\n`,
      )

      return {
        type: "tool_result" as const,
        tool_use_id: block.id,
        content: result,
      }
    }),
  )

  return results
}

function baseParams(systemPrompt: string, conversationMessages: Anthropic.MessageParam[]) {
  return {
    model: "claude-opus-4-6" as const,
    max_tokens: 32_000,
    system: systemPrompt,
    messages: conversationMessages,
    tools: anthropicToolDefinitions as Anthropic.Tool[],
    thinking: { type: "adaptive" as const },
  }
}

export async function POST(request: NextRequest): Promise<Response> {
  const { auth0 } = await import("@/lib/auth0")
  const session = await auth0.getSession()
  if (!session?.user?.sub) {
    return Response.json({ error: "Unauthorized" }, { status: 401 })
  }

  let client: Anthropic
  try {
    client = getClient()
  } catch (err) {
    return Response.json(
      { error: err instanceof Error ? err.message : "Configuration error" },
      { status: 500 },
    )
  }

  let body: { traceId: string; messages: ChatMessage[] }
  try {
    body = await request.json()
  } catch {
    return Response.json({ error: "Invalid JSON" }, { status: 400 })
  }

  const { traceId, messages } = body
  if (!traceId || !messages?.length) {
    return Response.json({ error: "traceId and messages are required" }, { status: 400 })
  }

  const tracePrefix = traceId.substring(0, 33)
  const traceData = await getTraceData(tracePrefix)
  if (!traceData || (traceData.rawLlmCalls.length === 0 && traceData.errors.length === 0)) {
    return Response.json({ error: "Trace not found" }, { status: 404 })
  }

  const indexed = indexTraceData(traceData)
  const systemPrompt = buildSummaryPrompt(indexed)
  const conversationMessages: Anthropic.MessageParam[] = messages.map(m => ({
    role: m.role,
    content: m.content,
  }))

  const encoder = new TextEncoder()

  const stream = new ReadableStream({
    async start(controller) {
      const enqueue = (data: string) => controller.enqueue(encoder.encode(data))
      const keepalive = setInterval(() => enqueue(": keepalive\n\n"), KEEPALIVE_INTERVAL_MS)

      try {
        let toolRounds = 0
        let emittedText = false

        while (true) {
          enqueue(
            `data: ${JSON.stringify({ type: "status", message: toolRounds === 0 ? "Thinking…" : "Analyzing…" })}\n\n`,
          )

          // Redact thinking blocks from prior rounds (each can be 10-50KB)
          for (const msg of conversationMessages) {
            if (msg.role !== "assistant" || !Array.isArray(msg.content)) continue
            for (const block of msg.content) {
              if ((block as { type: string }).type === "thinking") {
                ;(block as { thinking: string }).thinking = ""
              }
            }
          }

          const messageStream = client.messages.stream(
            baseParams(systemPrompt, conversationMessages),
          )
          const timeout = setTimeout(() => messageStream.abort(), ROUND_TIMEOUT_MS)
          let response: Anthropic.Message
          try {
            response = await messageStream.finalMessage()
          } finally {
            clearTimeout(timeout)
          }

          if (response.stop_reason !== "tool_use") {
            for (const block of response.content) {
              if (block.type === "text" && block.text) {
                enqueue(`data: ${JSON.stringify({ type: "text", text: block.text })}\n\n`)
                emittedText = true
              }
            }
            break
          }

          conversationMessages.push({ role: "assistant", content: response.content })
          const toolResults = await processToolCalls(response.content, indexed, enqueue)
          conversationMessages.push({ role: "user", content: toolResults })
          toolRounds++
        }

        // Force a text response if the model only produced thinking blocks
        if (!emittedText) {
          enqueue(`data: ${JSON.stringify({ type: "status", message: "Summarizing…" })}\n\n`)

          for (const msg of conversationMessages) {
            if (msg.role !== "assistant" || !Array.isArray(msg.content)) continue
            for (const block of msg.content) {
              if ((block as { type: string }).type === "thinking") {
                ;(block as { thinking: string }).thinking = ""
              }
            }
          }

          const { tools: _, ...noToolsParams } = baseParams(systemPrompt, conversationMessages)
          const finalStream = client.messages.stream(noToolsParams)
          const timeout = setTimeout(() => finalStream.abort(), ROUND_TIMEOUT_MS)
          try {
            const finalResponse = await finalStream.finalMessage()
            for (const block of finalResponse.content) {
              if (block.type === "text") {
                enqueue(`data: ${JSON.stringify({ type: "text", text: block.text })}\n\n`)
              }
            }
          } finally {
            clearTimeout(timeout)
          }
        }

        enqueue("data: [DONE]\n\n")
      } catch (err) {
        const message =
          err instanceof Anthropic.APIError
            ? `API error: ${err.status} ${err.message}`
            : err instanceof Error
              ? err.message
              : "Stream error"
        enqueue(`data: ${JSON.stringify({ error: message })}\n\n`)
      } finally {
        clearInterval(keepalive)
        controller.close()
      }
    },
  })

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