import type { LLMTool } from "@lmstudio/sdk";
import { createHash, randomUUID } from "node:crypto";
import { Validator, type Schema } from "jsonschema";
import type { Message } from "./messages";

export type CompletedTool = { id: string; name: string; arguments: string; parsed: Record<string, unknown> };

export class ToolBridge {
  readonly tools: LLMTool[];
  private readonly aliases = new Map<string, string>();
  private readonly originals = new Map<string, LLMTool>();
  private readonly normalized = new Map<string, string | null>();
  private readonly validator = new Validator();

  constructor(definitions: LLMTool[]) {
    const reserved = new Set(definitions.map(tool => tool.function.name));
    const normalize = (name: string) => name.replace(/[^A-Za-z0-9_-]+/g, "_");
    const counts = new Map<string, number>();
    for (const tool of definitions) {
      const key = normalize(tool.function.name);
      counts.set(key, (counts.get(key) ?? 0) + 1);
      this.normalized.set(key, counts.get(key) === 1 ? tool.function.name : null);
    }
    const used = new Set<string>();
    for (const tool of definitions) {
      const name = tool.function.name;
      if (this.originals.has(name)) throw new Error(`Duplicate tool name: ${name}`);
      this.originals.set(name, tool);
      let alias = name;
      if (!/^[A-Za-z0-9_-]{1,64}$/.test(name)) {
        const slug = name.replace(/[^A-Za-z0-9_-]+/g, "_").slice(0, 40) || "tool";
        const simple = normalize(name);
        if (simple.length <= 64 && counts.get(simple) === 1 && !reserved.has(simple)) alias = simple;
        else {
          let salt = 0;
          do {
            const hash = createHash("sha256").update(`${name}:${salt++}`).digest("hex").slice(0, 16);
            alias = `${slug}_${hash}`;
          } while (reserved.has(alias) || used.has(alias));
        }
      }
      used.add(alias);
      this.aliases.set(name, alias);
    }
    const names = [...this.aliases.keys()].sort((a, b) => b.length - a.length);
    const rewrite = (value: unknown): unknown => {
      if (typeof value === "string") {
        // One pass avoids rewriting an alias a second time when names overlap.
        const pattern = names.map(name => name.replace(/[.*+?^${}()|[\]\\]/g, "\\$&")).join("|");
        return pattern ? value.replace(new RegExp(pattern, "g"), name => this.aliases.get(name)!) : value;
      }
      if (Array.isArray(value)) return value.map(rewrite);
      if (value && typeof value === "object") return Object.fromEntries(Object.entries(value).map(([key, item]) => [key, rewrite(item)]));
      return value;
    };
    this.tools = definitions.map(tool => ({ type: "function", function: {
      name: this.aliases.get(tool.function.name)!,
      ...(tool.function.description ? { description: rewrite(tool.function.description) as string } : {}),
      // Keep schema keys, enums and defaults exact. Only rewrite natural-language descriptions.
      parameters: this.rewriteSchemaDescriptions(tool.function.parameters ?? {}, rewrite),
    } }));
  }

  private rewriteSchemaDescriptions(value: any, rewrite: (value: unknown) => unknown): any {
    if (Array.isArray(value)) return value.map(item => this.rewriteSchemaDescriptions(item, rewrite));
    if (value && typeof value === "object") return Object.fromEntries(Object.entries(value).map(([key, item]) => [key,
      ["enum", "const", "default", "examples"].includes(key) ? item :
      key === "description" && typeof item === "string" ? rewrite(item) : this.rewriteSchemaDescriptions(item, rewrite)]));
    return value;
  }

  original(name: string): string | undefined {
    if (this.originals.has(name)) return name;
    for (const [original, alias] of this.aliases) if (alias === name) return original;
    // Models sometimes normalize spaces themselves or omit a collision suffix.
    // Only accept this spelling if exactly one currently declared tool matches.
    return this.normalized.get(name) ?? undefined;
  }

  mapHistory(messages: Message[]): Message[] {
    return messages.map(message => {
      if (!message.tool_calls) return message;
      return { ...message, tool_calls: message.tool_calls.map((entry: any) => ({ ...entry,
        function: { ...entry.function, name: this.aliases.get(entry.function.name) ?? entry.function.name },
      })) };
    });
  }

  // Only an entire answer consisting of complete tool envelopes qualifies.
  // Quoted examples, fenced code, prose, unknown tools and invalid arguments stay text.
  recoverXml(text: string): CompletedTool[] | undefined {
    const result: CompletedTool[] = [];
    const outer = /\s*<tool_call>\s*<function=([^<>\r\n]+)>\s*([\s\S]*?)\s*<\/function>\s*<\/tool_call>\s*/gy;
    let at = 0;
    while (at < text.length) {
      outer.lastIndex = at;
      const match = outer.exec(text);
      if (!match) return;
      at = outer.lastIndex;
      const name = this.original(match[1].trim());
      if (!name || result.length >= 32) return;
      const schema = this.originals.get(name)!.function.parameters as Schema | undefined;
      const args: Record<string, unknown> = Object.create(null);
      const parameters = /\s*<parameter=([A-Za-z0-9_-]+)>\s*([\s\S]*?)\s*<\/parameter>\s*/gy;
      let offset = 0;
      while (offset < match[2].length) {
        parameters.lastIndex = offset;
        const parameter = parameters.exec(match[2]);
        if (!parameter) return;
        offset = parameters.lastIndex;
        const key = parameter[1];
        if (Object.hasOwn(args, key) || !Object.hasOwn(schema?.properties ?? {}, key)) return;
        const property = schema!.properties![key];
        const types = typeof property === "object" ? [property.type].flat() : [];
        const value = parameter[2].trim();
        if (types.includes("string")) args[key] = value;
        else {
          try { args[key] = JSON.parse(value); } catch { return; }
        }
      }
      try { if (schema && !this.validator.validate(args, schema).valid) return; }
      catch { return; }
      result.push({ id: `call_${randomUUID().replace(/-/g, "")}`, name, arguments: JSON.stringify(args), parsed: args });
    }
    return result.length ? result : undefined;
  }
}

import type { LLMTool } from "@lmstudio/sdk";
import { createHash, randomUUID } from "node:crypto";
import { Validator, type Schema } from "jsonschema";
import type { Message } from "./messages";

export type CompletedTool = { id: string; name: string; arguments: string; parsed: Record<string, unknown> };

export class ToolBridge {
  readonly tools: LLMTool[];
  private readonly aliases = new Map<string, string>();
  private readonly originals = new Map<string, LLMTool>();
  private readonly normalized = new Map<string, string | null>();
  private readonly validator = new Validator();

  constructor(definitions: LLMTool[]) {
    const reserved = new Set(definitions.map(tool => tool.function.name));
    const normalize = (name: string) => name.replace(/[^A-Za-z0-9_-]+/g, "_");
    const counts = new Map<string, number>();
    for (const tool of definitions) {
      const key = normalize(tool.function.name);
      counts.set(key, (counts.get(key) ?? 0) + 1);
      this.normalized.set(key, counts.get(key) === 1 ? tool.function.name : null);
    }
    const used = new Set<string>();
    for (const tool of definitions) {
      const name = tool.function.name;
      if (this.originals.has(name)) throw new Error(`Duplicate tool name: ${name}`);
      this.originals.set(name, tool);
      let alias = name;
      if (!/^[A-Za-z0-9_-]{1,64}$/.test(name)) {
        const slug = name.replace(/[^A-Za-z0-9_-]+/g, "_").slice(0, 40) || "tool";
        const simple = normalize(name);
        if (simple.length <= 64 && counts.get(simple) === 1 && !reserved.has(simple)) alias = simple;
        else {
          let salt = 0;
          do {
            const hash = createHash("sha256").update(`${name}:${salt++}`).digest("hex").slice(0, 16);
            alias = `${slug}_${hash}`;
          } while (reserved.has(alias) || used.has(alias));
        }
      }
      used.add(alias);
      this.aliases.set(name, alias);
    }
    const names = [...this.aliases.keys()].sort((a, b) => b.length - a.length);
    const rewrite = (value: unknown): unknown => {
      if (typeof value === "string") {
        // One pass avoids rewriting an alias a second time when names overlap.
        const pattern = names.map(name => name.replace(/[.*+?^${}()|[\]\\]/g, "\\$&")).join("|");
        return pattern ? value.replace(new RegExp(pattern, "g"), name => this.aliases.get(name)!) : value;
      }
      if (Array.isArray(value)) return value.map(rewrite);
      if (value && typeof value === "object") return Object.fromEntries(Object.entries(value).map(([key, item]) => [key, rewrite(item)]));
      return value;
    };
    this.tools = definitions.map(tool => ({ type: "function", function: {
      name: this.aliases.get(tool.function.name)!,
      ...(tool.function.description ? { description: rewrite(tool.function.description) as string } : {}),
      // Keep schema keys, enums and defaults exact. Only rewrite natural-language descriptions.
      parameters: this.rewriteSchemaDescriptions(tool.function.parameters ?? {}, rewrite),
    } }));
  }

  private rewriteSchemaDescriptions(value: any, rewrite: (value: unknown) => unknown): any {
    if (Array.isArray(value)) return value.map(item => this.rewriteSchemaDescriptions(item, rewrite));
    if (value && typeof value === "object") return Object.fromEntries(Object.entries(value).map(([key, item]) => [key,
      ["enum", "const", "default", "examples"].includes(key) ? item :
      key === "description" && typeof item === "string" ? rewrite(item) : this.rewriteSchemaDescriptions(item, rewrite)]));
    return value;
  }

  original(name: string): string | undefined {
    if (this.originals.has(name)) return name;
    for (const [original, alias] of this.aliases) if (alias === name) return original;
    // Models sometimes normalize spaces themselves or omit a collision suffix.
    // Only accept this spelling if exactly one currently declared tool matches.
    return this.normalized.get(name) ?? undefined;
  }

  mapHistory(messages: Message[]): Message[] {
    return messages.map(message => {
      if (!message.tool_calls) return message;
      return { ...message, tool_calls: message.tool_calls.map((entry: any) => ({ ...entry,
        function: { ...entry.function, name: this.aliases.get(entry.function.name) ?? entry.function.name },
      })) };
    });
  }

  // Only an entire answer consisting of complete tool envelopes qualifies.
  // Quoted examples, fenced code, prose, unknown tools and invalid arguments stay text.
  recoverXml(text: string): CompletedTool[] | undefined {
    const result: CompletedTool[] = [];
    const outer = /\s*<tool_call>\s*<function=([^<>\r\n]+)>\s*([\s\S]*?)\s*<\/function>\s*<\/tool_call>\s*/gy;
    let at = 0;
    while (at < text.length) {
      outer.lastIndex = at;
      const match = outer.exec(text);
      if (!match) return;
      at = outer.lastIndex;
      const name = this.original(match[1].trim());
      if (!name || result.length >= 32) return;
      const schema = this.originals.get(name)!.function.parameters as Schema | undefined;
      const args: Record<string, unknown> = Object.create(null);
      const parameters = /\s*<parameter=([A-Za-z0-9_-]+)>\s*([\s\S]*?)\s*<\/parameter>\s*/gy;
      let offset = 0;
      while (offset < match[2].length) {
        parameters.lastIndex = offset;
        const parameter = parameters.exec(match[2]);
        if (!parameter) return;
        offset = parameters.lastIndex;
        const key = parameter[1];
        if (Object.hasOwn(args, key) || !Object.hasOwn(schema?.properties ?? {}, key)) return;
        const property = schema!.properties![key];
        const types = typeof property === "object" ? [property.type].flat() : [];
        const value = parameter[2].trim();
        if (types.includes("string")) args[key] = value;
        else {
          try { args[key] = JSON.parse(value); } catch { return; }
        }
      }
      try { if (schema && !this.validator.validate(args, schema).valid) return; }
      catch { return; }
      result.push({ id: `call_${randomUUID().replace(/-/g, "")}`, name, arguments: JSON.stringify(args), parsed: args });
    }
    return result.length ? result : undefined;
  }
}