import type Anthropic from "@anthropic-ai/sdk";
import { t, toJSONSchema } from "structural";
import type {
  Compiler,
  CompilerImplementationParams,
  CompilerIR,
  CompilerModalities,
} from "./compiler-interface.ts";
import { compilerUsage, defineCompiler } from "./compiler-interface.ts";
import { parseToolCall } from "./parse-tool-call.ts";
import type {
  Agent,
  AnthropicAssistantData,
  AssistantMessage,
  Content as IRContent,
  MalformedToolRequest,
} from "../llm-ir.ts";
import type { LoadedTools, ToolCall } from "../tool-def.ts";
import { errorToString, err } from "../result.ts";
import * as irPrompts from "./ir-prompts.ts";

type ToolCallRequest<A extends Agent<any, any, any>> = ToolCall<A["tools"]>;
type LoadedTool<A extends Agent<any, any, any>> = LoadedTools<A["tools"]>[keyof LoadedTools<
  A["tools"]
>];

export type AnthropicCompilerModel = {
  client: Anthropic;
  model: string;
  maxTokens: number;
  modalities?: CompilerModalities;
  thinking?: Anthropic.ThinkingConfigParam;
  outputConfig?: Anthropic.OutputConfig;
};

const ThinkingBlockSchema = t.subtype({
  type: t.value("thinking"),
  thinking: t.str,
  signature: t.str,
});

function imagePlaceholderContent(): string {
  return irPrompts.imageAttachmentPlaceholderText();
}

function anthropicContentParts(
  content: IRContent["content"],
  modalities?: CompilerModalities,
): Array<Anthropic.TextBlockParam | Anthropic.ImageBlockParam> {
  const output: Array<Anthropic.TextBlockParam | Anthropic.ImageBlockParam> = [];
  for (const part of content) {
    if (part.type === "text") {
      output.push({ type: "text", text: part.content });
      continue;
    }
    if (modalities?.includes("vision")) {
      output.push({
        type: "image",
        source: {
          type: "base64",
          media_type: part.image.mimeType,
          data: part.image.base64Data,
        },
      });
    } else {
      output.push({ type: "text", text: imagePlaceholderContent() });
    }
  }
  return output;
}

function toModelMessage<A extends Agent<any, any, any>>(
  messages: Array<CompilerIR<A>>,
  modalities?: CompilerModalities,
): Array<Anthropic.MessageParam> {
  const output: Anthropic.MessageParam[] = [];

  for (const ir of messages) {
    output.push(modelMessageFromIr(ir, modalities));
  }

  return output;
}

function modelMessageFromIr<A extends Agent<any, any, any>>(
  ir: CompilerIR<A>,
  modalities?: CompilerModalities,
): Anthropic.MessageParam {
  if (ir.role === "assistant") {
    let thinkingBlocks = ir.anthropic?.thinkingBlocks || [];
    const toolCalls = ir.toolCalls || [];
    return {
      role: "assistant",
      content: [
        ...thinkingBlocks,
        { type: "text", text: ir.content || " " },
        ...toolCalls
          .filter((t: any) => t.type === "tool-call")
          .map((t: any) => {
            return {
              type: "tool_use" as const,
              id: t.toolCallId,
              name: t.name,
              input: t.original || {},
            };
          }),
      ],
    };
  }

  if (ir.role === "user") {
    return {
      role: "user",
      content: anthropicContentParts(ir.content, modalities),
    };
  }

  if (ir.role === "tool-output") {
    return {
      role: "user",
      content: [
        {
          type: "tool_result",
          tool_use_id: ir.toolCall.toolCallId,
          content: anthropicContentParts(ir.content, modalities),
        },
      ],
    };
  }

  if (ir.role === "tool-skip-output") {
    return {
      role: "user",
      content: [
        {
          type: "tool_result",
          tool_use_id: ir.toolCall.toolCallId,
          is_error: true,
          content: irPrompts.toolSkip(ir.reason),
        },
      ],
    };
  }

  if (ir.role === "tool-runtime-error") {
    return {
      role: "user",
      content: [
        {
          type: "tool_result",
          tool_use_id: ir.toolCall.toolCallId,
          is_error: true,
          content: `Error: ${ir.error}`,
        },
      ],
    };
  }

  if (ir.role === "tool-parse-error") {
    return {
      role: "user",
      content: [
        {
          type: "tool_result",
          tool_use_id: ir.malformedRequest.toolCallId,
          is_error: true,
          content: `Error: ${ir.malformedRequest.error}`,
        },
      ],
    };
  }

  if (ir.role === "tool-validation-error") {
    return {
      role: "user",
      content: [
        {
          type: "tool_result",
          tool_use_id: ir.toolCall.toolCallId,
          is_error: true,
          content: `Error: ${ir.error}`,
        },
      ],
    };
  }

  if (ir.role === "lowered-checkpoint") {
    return {
      role: "user",
      content: anthropicContentParts(ir.content, modalities),
    };
  }

  const _: never = ir;
  throw new Error(`Unsupported IR role: ${(ir as any).role}`);
}

function generateCurlFrom(params: {
  baseURL: string;
  model: string;
  system: string;
  messages: Array<Anthropic.MessageParam>;
  tools?: Array<{ description: string; input_schema: any; name: string }>;
  maxTokens: number;
  thinking?: Anthropic.ThinkingConfigParam;
  outputConfig?: Anthropic.OutputConfig;
}): string {
  const { baseURL, model, system, messages, tools, maxTokens, thinking, outputConfig } = params;
  const requestBody = {
    model,
    system,
    messages,
    tools,
    tool_choice: {
      type: "auto",
      disable_parallel_tool_use: false,
    },
    max_tokens: maxTokens,
    ...(thinking ? { thinking } : {}),
    ...(outputConfig ? { output_config: outputConfig } : {}),
    stream: true,
  };

  // Curl requests need an API Version
  // Currently hardcoded in Anthropic SDK
  const ANTHROPIC_API_VERSION = "2023-06-01";

  return `curl -X POST "${baseURL}/v1/messages" \\
  -H "Content-Type: application/json" \\
  -H "x-api-key: [REDACTED_API_KEY]" \\
  -H "anthropic-version: ${ANTHROPIC_API_VERSION}" \\
  -d @- <<'JSON'
${JSON.stringify(requestBody)}
JSON`;
}

export const runAnthropicAgent: Compiler<AnthropicCompilerModel> = defineCompiler(
  async <A extends Agent<any, any, any>>(
    params: CompilerImplementationParams<A, AnthropicCompilerModel>,
  ) => {
    const { model, irs, abortSignal, transport, systemPrompt, autofixJson } = params;
    const messages = toModelMessage(irs, model.modalities);
    const sysPrompt = systemPrompt ? await systemPrompt() : "";

    const toolDefs = params.tools || {};
    const toolDefinitions: Array<{ description: string; input_schema: any; name: string }> = [];
    const toolEntries = Object.entries(toolDefs) as Array<[string, LoadedTool<A>]>;
    toolEntries.forEach(([name, toolDef]) => {
      const argJsonSchema = toJSONSchema("ignore", toolDef.ArgumentsSchema);
      // Delete JSON schema fields unused by AI SDK
      // @ts-ignore
      delete argJsonSchema.$schema;
      delete argJsonSchema.description;
      // @ts-ignore
      delete argJsonSchema.title;

      toolDefinitions.push({
        name,
        description: toolDef.description,
        input_schema: argJsonSchema,
      });
    });
    const toolParams =
      toolDefinitions.length === 0
        ? {}
        : {
            tools: toolDefinitions,
          };

    const modelOptions = {
      ...(model.thinking ? { thinking: model.thinking } : {}),
      ...(model.outputConfig ? { output_config: model.outputConfig } : {}),
    };

    const curl = generateCurlFrom({
      baseURL: model.client.baseURL,
      model: model.model,
      system: sysPrompt,
      messages,
      ...toolParams,
      maxTokens: model.maxTokens,
      thinking: model.thinking,
      outputConfig: model.outputConfig,
    });

    try {
      const system = sysPrompt == null ? {} : { system: sysPrompt };
      const { data: stream, response } = await model.client.messages
        .create({
          ...system,
          model: model.model,
          messages,
          ...toolParams,
          tool_choice: {
            type: "auto",
            disable_parallel_tool_use: false,
          },
          max_tokens: model.maxTokens,
          ...modelOptions,
          stream: true,
        })
        .withResponse();

      let content = "";
      let reasoningContent: string | undefined = undefined;
      let usage = {
        input: 0,
        cachedInput: 0,
        output: 0,
      };
      const thinkingBlocks: Array<
        | {
            type: "thinking";
            thinking?: string;
            signature?: string;
            index: number;
          }
        | {
            type: "redacted_thinking";
            data: string;
          }
      > = [];
      let inProgressTools = new Map<
        number,
        {
          id: string;
          index: number;
          name: string;
          partialJson: string;
        }
      >();
      try {
        for await (const chunk of stream) {
          if (abortSignal.aborted) break;

          switch (chunk.type) {
            case "content_block_delta":
              switch (chunk.delta.type) {
                case "text_delta":
                  content += chunk.delta.text;
                  params.onTokens(chunk.delta.text, "content");
                  break;
                case "thinking_delta":
                  if (reasoningContent == null) reasoningContent = "";
                  reasoningContent += chunk.delta.thinking;
                  params.onTokens(chunk.delta.thinking, "reasoning");
                  if (thinkingBlocks.length === 0) {
                    thinkingBlocks.push({
                      type: "thinking",
                      thinking: chunk.delta.thinking,
                      index: chunk.index,
                    });
                  } else {
                    const lastBlock = thinkingBlocks[thinkingBlocks.length - 1];
                    if (lastBlock.type === "thinking" && lastBlock.index === chunk.index) {
                      lastBlock.thinking += chunk.delta.thinking;
                    } else {
                      thinkingBlocks.push({
                        type: "thinking",
                        thinking: chunk.delta.thinking,
                        index: chunk.index,
                      });
                    }
                  }
                  break;
                case "signature_delta":
                  if (thinkingBlocks.length === 0) {
                    thinkingBlocks.push({
                      type: "thinking",
                      signature: chunk.delta.signature,
                      index: chunk.index,
                    });
                  } else {
                    const lastBlock = thinkingBlocks[thinkingBlocks.length - 1];
                    if (lastBlock.type === "thinking" && lastBlock.index === chunk.index) {
                      lastBlock.signature = chunk.delta.signature;
                    } else {
                      thinkingBlocks.push({
                        type: "thinking",
                        signature: chunk.delta.signature,
                        index: chunk.index,
                      });
                    }
                  }
                  break;
                case "input_json_delta":
                  {
                    const tool = inProgressTools.get(chunk.index);
                    if (tool != null) {
                      params.onTokens(chunk.delta.partial_json, "tool");
                      tool.partialJson += chunk.delta.partial_json;
                    }
                  }
                  break;
              }
              break;
            case "content_block_start":
              switch (chunk.content_block.type) {
                case "tool_use":
                  params.onTokens(chunk.content_block.name, "tool");
                  inProgressTools.set(chunk.index, {
                    id: chunk.content_block.id,
                    index: chunk.index,
                    name: chunk.content_block.name,
                    partialJson: "",
                  });
                  break;
                case "redacted_thinking":
                  thinkingBlocks.push({
                    type: "redacted_thinking",
                    data: chunk.content_block.data,
                  });
                  break;
              }
              break;

            case "message_delta":
              usage.output = chunk.usage.output_tokens;
              if (chunk.usage.input_tokens && chunk.usage.input_tokens > 0) {
                usage.input = chunk.usage.input_tokens;
              }
              usage.cachedInput = chunk.usage.cache_read_input_tokens ?? usage.cachedInput;
              break;
            case "message_start":
              usage.input = chunk.message.usage.input_tokens;
              usage.cachedInput = chunk.message.usage.cache_read_input_tokens ?? 0;
              break;
          }
        }
      } catch (e) {
        if (!abortSignal.aborted) {
          return err({
            type: "stream-error",
            requestError: errorToString(e),
            curl,
            usage: compilerUsage(usage.input, usage.output, usage.cachedInput),
            headers: response.headers,
          });
        }
      }

      const compilerTokens = compilerUsage(usage.input, usage.output, usage.cachedInput);

      let anthropic: { anthropic?: AnthropicAssistantData } = {};
      if (thinkingBlocks.length > 0) {
        anthropic.anthropic = {
          thinkingBlocks: thinkingBlocks.map(b => {
            if (b.type === "redacted_thinking") return b;
            return ThinkingBlockSchema.slice({
              type: "thinking",
              signature: b.signature || "",
              thinking: b.thinking || "",
            });
          }),
        };
      }

      const assistantMessage: AssistantMessage<A["tools"]> = {
        role: "assistant",
        content,
        reasoningContent,
        ...anthropic,
        usage: compilerTokens,
      };

      return params.finish({
        curl,
        headers: response.headers,
        usage: compilerTokens,
        abortedOutput: assistantMessage,
        parsedOutput: async () => {
          if (inProgressTools.size === 0) return assistantMessage;

          // Sort tool calls by their content block index to preserve ordering
          const sortedTools = Array.from(inProgressTools.entries())
            .sort(([a], [b]) => a - b)
            .map(([_, v]) => v);

          const toolCalls: Array<ToolCallRequest<A> | MalformedToolRequest> = [];

          for (const inProgressTool of sortedTools) {
            const chatToolCall = {
              toolCallId: inProgressTool.id,
              toolName: inProgressTool.name,
              args: inProgressTool.partialJson,
            };
            const parseResult = await parseToolCall<A["tools"]>({
              toolCall: chatToolCall,
              toolDefs,
              autofixJson,
              abortSignal,
              transport,
            });

            if (parseResult.status === "error") {
              toolCalls.push({
                type: "malformed-tool-request",
                error: parseResult.message,
                toolCallId: inProgressTool.id,
                call: {
                  original: {
                    name: inProgressTool.name,
                    arguments: inProgressTool.partialJson,
                  },
                },
              });
              continue;
            }

            toolCalls.push(parseResult.tool);
          }

          if (toolCalls.length > 0) assistantMessage.toolCalls = toolCalls;
          return assistantMessage;
        },
      });
    } catch (e) {
      return err({
        type: "request-error",
        requestError: errorToString(e),
        curl,
      });
    }
  },
);