import { Emitter } from "@noya-app/emitter";
import { encodeSchema } from "@noya-app/noya-schemas";
import { Base64, isDeepEqual } from "@noya-app/noya-utils";
import { Observable } from "@noya-app/observable";
import type { Static, TObject } from "@sinclair/typebox";
import { Value } from "@sinclair/typebox/value";
import type {
  AIGenerateImage,
  AIGenerateImageRequest,
  AIGenerateImageResponse,
  AIGenerateRequest,
  AIGenerateResponse,
  AIGenerateTextRequest,
  AIGenerateTextResponse,
} from "./rpc/routes";
import { RPCManager } from "./rpcManager";

export type AIToolDefinition = {
  functionName: string;
  description: string;
  parameters: Record<string, unknown>;
  onCall: (parameters: Record<string, unknown>) => Promise<string>;
};

export type CallableAIToolsMap = Record<
  string,
  { description: string; parameters: Record<string, unknown> }
>;

export type AIToolInvocation = {
  id: string;
  name: string;
  parameters: Record<string, unknown>;
};

export type AIImageInput =
  | Blob
  | { data: Uint8Array; mediaType: string }
  | { url: string | URL };

export type AIGenerateObjectOptions<Schema extends TObject> = {
  prompt: string;
  schema: Schema;
  system?: string;
  images?: AIImageInput[];
  model?: string;
};

export type AIGenerateTextOptions = {
  prompt: string;
  system?: string;
  images?: AIImageInput[];
  model?: string;
};

export type AIGenerateImageOptions = {
  prompt: string;
  model?: string;
};

export type AIGeneratedImage = {
  data: Uint8Array;
  mediaType: string;
};

export class AIManager {
  invocationEmitter = new Emitter<[AIToolInvocation]>();
  responseEmitter = new Emitter<[AIToolInvocation, string | undefined]>();

  tools$ = new Observable<AIToolDefinition[]>([]);
  systemMessage$ = new Observable<string | undefined>(undefined);

  callableTools$ = this.tools$.map(
    (tools): CallableAIToolsMap =>
      Object.fromEntries(
        tools.map(({ functionName, description, parameters }) => [
          functionName,
          { description, parameters },
        ])
      )
  );
  configuration$ = Observable.combine(
    [this.callableTools$, this.systemMessage$],
    ([callableTools, systemMessage]) => ({
      tools: callableTools,
      systemMessage,
    }),
    { isEqual: isDeepEqual }
  );

  get callableTools() {
    return this.callableTools$.get();
  }

  constructor(public rpcManager: RPCManager) {
    this.fetch = this.createFetch();
  }

  private createFetch(): typeof globalThis.fetch {
    return (async (_input: RequestInfo | URL, init?: RequestInit) => {
      // Create a TransformStream to handle the response body
      const { readable, writable } = new TransformStream();
      const writer = writable.getWriter();

      // Make the RPC request with streaming
      this.rpcManager
        .requestStreamingRoute("POST /api/ai", {
          body: init?.body?.toString(),
          headers: init?.headers
            ? Object.fromEntries(new Headers(init.headers).entries())
            : undefined,
          onStreamChunk: async (chunk) => {
            // Write each chunk to the stream
            await writer.write(new TextEncoder().encode(chunk));
          },
          onStreamEnd: async () => {
            await writer.close();
          },
        })
        .then(
          async () => {
            // Nothing to do here - writer is closed in onStreamEnd
          },
          async (error) => {
            await writer.abort(error);
          }
        );

      // Create response object immediately with default values
      const response = new Response(readable, {
        status: 200,
        statusText: "OK",
        headers: new Headers({
          "content-type": "text/plain; charset=utf-8",
        }),
      });

      return response;
    }) as typeof globalThis.fetch;
  }

  fetch: typeof globalThis.fetch;

  registerTool(tool: AIToolDefinition) {
    this.tools$.set([...this.tools$.get(), tool]);

    return () => this.unregisterTool(tool.functionName);
  }

  unregisterTool(name: string) {
    this.tools$.set(
      this.tools$.get().filter((tool) => tool.functionName !== name)
    );
  }

  callTool(toolInvocation: AIToolInvocation) {
    const tool = this.tools$
      .get()
      .find((tool) => tool.functionName === toolInvocation.name);

    if (!tool) {
      console.error(`Tool ${toolInvocation.name} not found`);
      return;
    }

    return tool.onCall(toolInvocation.parameters);
  }

  async generateObject<Schema extends TObject>({
    prompt,
    schema,
    system,
    images = [],
    model,
  }: AIGenerateObjectOptions<Schema>): Promise<Static<Schema>> {
    validatePromptAndModel(prompt, model);

    const payload: AIGenerateRequest = {
      prompt,
      schema,
      encodedSchema: encodeSchema(schema),
      ...(system ? { system } : {}),
      ...(model ? { model } : {}),
      ...(images.length
        ? { images: await Promise.all(images.map(normalizeImageInput)) }
        : {}),
    };

    const response = await this.rpcManager.requestRoute(
      "POST /api/ai/generate",
      {
        body: JSON.stringify(payload),
        headers: { "Content-Type": "application/json" },
      }
    );
    const result = this.rpcManager.getResponseBody(response) as
      | AIGenerateResponse
      | undefined;

    if (!result || !("object" in result)) {
      throw new Error("AI returned an unexpected response");
    }

    if (!Value.Check(schema, result.object)) {
      const errors = [...Value.Errors(schema, result.object)]
        .map(({ path, message }) => `${path || "/"}: ${message}`)
        .join(", ");

      throw new Error(
        `AI response does not match the requested schema: ${errors}`
      );
    }

    return result.object as Static<Schema>;
  }

  async generateText({
    prompt,
    system,
    images = [],
    model,
  }: AIGenerateTextOptions): Promise<string> {
    validatePromptAndModel(prompt, model);

    const payload: AIGenerateTextRequest = {
      prompt,
      ...(system ? { system } : {}),
      ...(model ? { model } : {}),
      ...(images.length
        ? { images: await Promise.all(images.map(normalizeImageInput)) }
        : {}),
    };
    const response = await this.rpcManager.requestRoute(
      "POST /api/ai/generate/text",
      {
        body: JSON.stringify(payload),
        headers: { "Content-Type": "application/json" },
      }
    );
    const result = this.rpcManager.getResponseBody(response) as
      | AIGenerateTextResponse
      | undefined;

    if (!result || typeof result.text !== "string") {
      throw new Error("AI returned an unexpected text response");
    }

    return result.text;
  }

  async generateImage({
    prompt,
    model,
  }: AIGenerateImageOptions): Promise<AIGeneratedImage> {
    validatePromptAndModel(prompt, model);

    const payload: AIGenerateImageRequest = {
      prompt,
      ...(model ? { model } : {}),
    };
    const response = await this.rpcManager.requestRoute(
      "POST /api/ai/generate/image",
      {
        body: JSON.stringify(payload),
        headers: { "Content-Type": "application/json" },
      }
    );
    const result = this.rpcManager.getResponseBody(response) as
      | AIGenerateImageResponse
      | undefined;

    if (
      !result ||
      typeof result.base64 !== "string" ||
      typeof result.mediaType !== "string"
    ) {
      throw new Error("AI returned an unexpected image response");
    }

    return {
      data: Base64.decode(result.base64),
      mediaType: result.mediaType,
    };
  }
}

function validatePromptAndModel(prompt: string, model?: string) {
  if (!prompt.trim()) {
    throw new Error("AI prompt must not be empty");
  }
  if (model !== undefined && !model.trim()) {
    throw new Error("AI model must not be empty");
  }
}

async function normalizeImageInput(
  image: AIImageInput
): Promise<AIGenerateImage> {
  if (image instanceof Blob) {
    assertImageMediaType(image.type);

    return {
      type: "data",
      data: Base64.encode(await image.arrayBuffer()),
      mediaType: image.type,
    };
  }

  if ("url" in image) {
    const url = new URL(image.url).toString();

    if (!["http:", "https:"].includes(new URL(url).protocol)) {
      throw new Error("AI image URLs must use http or https");
    }

    return { type: "url", url };
  }

  assertImageMediaType(image.mediaType);

  return {
    type: "data",
    data: Base64.encode(image.data),
    mediaType: image.mediaType,
  };
}

function assertImageMediaType(mediaType: string) {
  if (!mediaType.toLowerCase().startsWith("image/")) {
    throw new Error(`Invalid AI image media type: ${mediaType || "(missing)"}`);
  }
}