import java.util.List;
import java.util.Optional;

import io.jooby.annotation.mcp.McpCompletion;
import io.jooby.annotation.mcp.McpPrompt;
import io.jooby.annotation.mcp.McpResource;
import io.jooby.annotation.mcp.McpTool;
import io.modelcontextprotocol.server.McpSyncServerExchange;

/** A collection of tools, prompts, and resources exposed to the LLM via MCP. */
public class CalculatorTools {

  /**
   * Adds two integers together and returns the result.
   *
   * @param a The first number to add.
   * @param b The second number to add.
   * @return The sum of the two numbers.
   */
  @McpTool(name = "add_numbers")
  public int add(int a, int b) {
    return a + b;
  }

  @McpPrompt(name = "math_tutor", description = "A prompt to initiate a math tutoring session")
  public String mathTutor(String topic) {
    return "You are a helpful math tutor. Please explain the concept of "
        + topic
        + " step by step.";
  }

  @McpResource(
      uri = "calculator://manual/usage",
      name = "Calculator Manual",
      description = "Instructions on how to use the calculator")
  public String manual() {
    return "The Calculator supports basic arithmetic. Use the add_numbers tool to sum values.";
  }

  @McpResource(
      uri = "calculator://history/{user}",
      name = "User History",
      description = "Retrieves the calculation history for a specific user")
  public String history(String user) {
    return "History for " + user + ":\n5 + 10 = 15\n2 * 4 = 8";
  }

  @McpCompletion(ref = "calculator://history/{user}")
  public List<String> historyCompletions(String user) {
    // In a real app, this would query a database for active usernames matching the input
    return List.of("alice", "bob", "charlie");
  }

  @McpTool(
      name = "get_session_info",
      description = "Returns the current MCP session ID using the injected exchange.")
  public String getSessionInfo(McpSyncServerExchange exchange) {
    return Optional.ofNullable(exchange.sessionId()).orElse("No active session");
  }
}

@io.jooby.annotation.Generated(CalculatorTools.class)
public class CalculatorToolsMcp_ implements io.jooby.mcp.McpService {
    protected java.util.function.Function<io.jooby.Context, CalculatorTools> factory;

    public CalculatorToolsMcp_() {
      this(io.jooby.SneakyThrows.singleton(CalculatorTools::new));
    }

    public CalculatorToolsMcp_(CalculatorTools instance) {
      setup(ctx -> instance);
    }

    public CalculatorToolsMcp_(io.jooby.SneakyThrows.Supplier<CalculatorTools> provider) {
      setup(ctx -> provider.get());
    }

    public CalculatorToolsMcp_(io.jooby.SneakyThrows.Function<Class<CalculatorTools>, CalculatorTools> provider) {
      setup(ctx -> provider.apply(CalculatorTools.class));
    }

    private void setup(java.util.function.Function<io.jooby.Context, CalculatorTools> factory) {
      this.factory = factory;
    }

    private io.modelcontextprotocol.json.McpJsonMapper json;
    @Override
    public void capabilities(io.modelcontextprotocol.spec.McpSchema.ServerCapabilities.Builder capabilities) {
      capabilities.tools(true);
      capabilities.prompts(true);
      capabilities.resources(true, true);
      capabilities.completions();
    }

    @Override
    public String serverKey() {
      return "default";
    }

    @Override
    public java.util.List<io.modelcontextprotocol.server.McpServerFeatures.SyncCompletionSpecification> completions() {
      var completions = new java.util.ArrayList<io.modelcontextprotocol.server.McpServerFeatures.SyncCompletionSpecification>();
      completions.add(new io.modelcontextprotocol.server.McpServerFeatures.SyncCompletionSpecification(new io.modelcontextprotocol.spec.McpSchema.PromptReference("math_tutor"), (exchange, req) -> new io.jooby.mcp.McpResult(this.json).toCompleteResult(java.util.List.of())));
      completions.add(new io.modelcontextprotocol.server.McpServerFeatures.SyncCompletionSpecification(new io.modelcontextprotocol.spec.McpSchema.ResourceReference("calculator://history/{user}"), (exchange, req) -> this.historyCompletionHandler(exchange, null, req)));
      return completions;
    }

    @Override
    public java.util.List<io.modelcontextprotocol.server.McpStatelessServerFeatures.SyncCompletionSpecification> statelessCompletions() {
      var completions = new java.util.ArrayList<io.modelcontextprotocol.server.McpStatelessServerFeatures.SyncCompletionSpecification>();
      completions.add(new io.modelcontextprotocol.server.McpStatelessServerFeatures.SyncCompletionSpecification(new io.modelcontextprotocol.spec.McpSchema.PromptReference("math_tutor"), (ctx, req) -> new io.jooby.mcp.McpResult(this.json).toCompleteResult(java.util.List.of())));
      completions.add(new io.modelcontextprotocol.server.McpStatelessServerFeatures.SyncCompletionSpecification(new io.modelcontextprotocol.spec.McpSchema.ResourceReference("calculator://history/{user}"), (ctx, req) -> this.historyCompletionHandler(null, ctx, req)));
      return completions;
    }

    @Override
    public void install(io.jooby.Jooby app, io.modelcontextprotocol.server.McpSyncServer server) throws Exception {
      this.json = app.require(io.modelcontextprotocol.json.McpJsonMapper.class);
      var schemaGenerator = app.require(com.github.victools.jsonschema.generator.SchemaGenerator.class);

      server.addTool(new io.modelcontextprotocol.server.McpServerFeatures.SyncToolSpecification(addToolSpec(schemaGenerator), (exchange, req) -> this.add(exchange, null, req)));
      server.addPrompt(new io.modelcontextprotocol.server.McpServerFeatures.SyncPromptSpecification(mathTutorPromptSpec(), (exchange, req) -> this.mathTutor(exchange, null, req)));
      server.addResource(new io.modelcontextprotocol.server.McpServerFeatures.SyncResourceSpecification(manualResourceSpec(), (exchange, req) -> this.manual(exchange, null, req)));
      server.addResourceTemplate(new io.modelcontextprotocol.server.McpServerFeatures.SyncResourceTemplateSpecification(historyResourceTemplateSpec(), (exchange, req) -> this.history(exchange, null, req)));
      server.addTool(new io.modelcontextprotocol.server.McpServerFeatures.SyncToolSpecification(getSessionInfoToolSpec(schemaGenerator), (exchange, req) -> this.getSessionInfo(exchange, null, req)));
    }

    @Override
    public void install(io.jooby.Jooby app, io.modelcontextprotocol.server.McpStatelessSyncServer server) throws Exception {
      this.json = app.require(io.modelcontextprotocol.json.McpJsonMapper.class);
      var schemaGenerator = app.require(com.github.victools.jsonschema.generator.SchemaGenerator.class);

      server.addTool(new io.modelcontextprotocol.server.McpStatelessServerFeatures.SyncToolSpecification(addToolSpec(schemaGenerator), (ctx, req) -> this.add(null, ctx, req)));
      server.addPrompt(new io.modelcontextprotocol.server.McpStatelessServerFeatures.SyncPromptSpecification(mathTutorPromptSpec(), (ctx, req) -> this.mathTutor(null, ctx, req)));
      server.addResource(new io.modelcontextprotocol.server.McpStatelessServerFeatures.SyncResourceSpecification(manualResourceSpec(), (ctx, req) -> this.manual(null, ctx, req)));
      server.addResourceTemplate(new io.modelcontextprotocol.server.McpStatelessServerFeatures.SyncResourceTemplateSpecification(historyResourceTemplateSpec(), (ctx, req) -> this.history(null, ctx, req)));
      server.addTool(new io.modelcontextprotocol.server.McpStatelessServerFeatures.SyncToolSpecification(getSessionInfoToolSpec(schemaGenerator), (ctx, req) -> this.getSessionInfo(null, ctx, req)));
    }

    private io.modelcontextprotocol.spec.McpSchema.Tool addToolSpec(com.github.victools.jsonschema.generator.SchemaGenerator schemaGenerator) {
      var schema = new java.util.LinkedHashMap<String, Object>();
      schema.put("type", "object");
      var props = new java.util.LinkedHashMap<String, Object>();
      schema.put("properties", props);
      var req = new java.util.ArrayList<String>();
      schema.put("required", req);
      var schema_a = schemaGenerator.generateSchema(int.class);
      props.put("a", schema_a);
      req.add("a");
      var schema_b = schemaGenerator.generateSchema(int.class);
      props.put("b", schema_b);
      req.add("b");
      return new io.modelcontextprotocol.spec.McpSchema.Tool("add_numbers", null, null, this.json.convertValue(schema, io.modelcontextprotocol.spec.McpSchema.JsonSchema.class), null, null, null);
    }

    private io.modelcontextprotocol.spec.McpSchema.CallToolResult add(io.modelcontextprotocol.server.McpSyncServerExchange exchange, io.modelcontextprotocol.common.McpTransportContext transportContext, io.modelcontextprotocol.spec.McpSchema.CallToolRequest req) {
      var ctx = exchange != null ? (io.jooby.Context) exchange.transportContext().get("CTX") : (transportContext != null ? (io.jooby.Context) transportContext.get("CTX") : null);
      var args = req.arguments() != null ? req.arguments() : java.util.Collections.<String, Object>emptyMap();
      var c = this.factory.apply(ctx);
      var raw_a = args.get("a");
      if (raw_a == null) throw new IllegalArgumentException("Missing req param: a");
      var a = ((Number) raw_a).intValue();
      var raw_b = args.get("b");
      if (raw_b == null) throw new IllegalArgumentException("Missing req param: b");
      var b = ((Number) raw_b).intValue();
      var result = c.add(a, b);
      return new io.jooby.mcp.McpResult(this.json).toCallToolResult(result, false);
    }

    private io.modelcontextprotocol.spec.McpSchema.Prompt mathTutorPromptSpec() {
      var args = new java.util.ArrayList<io.modelcontextprotocol.spec.McpSchema.PromptArgument>();
      args.add(new io.modelcontextprotocol.spec.McpSchema.PromptArgument("topic", null, false));
      return new io.modelcontextprotocol.spec.McpSchema.Prompt("math_tutor", null, "A prompt to initiate a math tutoring session", args);
    }

    private io.modelcontextprotocol.spec.McpSchema.GetPromptResult mathTutor(io.modelcontextprotocol.server.McpSyncServerExchange exchange, io.modelcontextprotocol.common.McpTransportContext transportContext, io.modelcontextprotocol.spec.McpSchema.GetPromptRequest req) {
      var ctx = exchange != null ? (io.jooby.Context) exchange.transportContext().get("CTX") : (transportContext != null ? (io.jooby.Context) transportContext.get("CTX") : null);
      var args = req.arguments() != null ? req.arguments() : java.util.Collections.<String, Object>emptyMap();
      var c = this.factory.apply(ctx);
      var raw_topic = args.get("topic");
      var topic = raw_topic != null ? raw_topic.toString() : null;
      var result = c.mathTutor(topic);
      return new io.jooby.mcp.McpResult(this.json).toPromptResult(result);
    }

    private io.modelcontextprotocol.spec.McpSchema.Resource manualResourceSpec() {
      return new io.modelcontextprotocol.spec.McpSchema.Resource("calculator://manual/usage", "Calculator Manual", null, "Instructions on how to use the calculator", io.jooby.MediaType.byFileExtension("calculator://manual/usage", "text/plain").getValue(), null, null, null);
    }

    private io.modelcontextprotocol.spec.McpSchema.ReadResourceResult manual(io.modelcontextprotocol.server.McpSyncServerExchange exchange, io.modelcontextprotocol.common.McpTransportContext transportContext, io.modelcontextprotocol.spec.McpSchema.ReadResourceRequest req) {
      var ctx = exchange != null ? (io.jooby.Context) exchange.transportContext().get("CTX") : (transportContext != null ? (io.jooby.Context) transportContext.get("CTX") : null);
      var args = java.util.Collections.<String, Object>emptyMap();
      var c = this.factory.apply(ctx);
      var result = c.manual();
      return new io.jooby.mcp.McpResult(this.json).toResourceResult(req.uri(), result);
    }

    private io.modelcontextprotocol.spec.McpSchema.ResourceTemplate historyResourceTemplateSpec() {
      return new io.modelcontextprotocol.spec.McpSchema.ResourceTemplate("calculator://history/{user}", "User History", null, "Retrieves the calculation history for a specific user", io.jooby.MediaType.byFileExtension("calculator://history/{user}", "text/plain").getValue(), null, null);
    }

    private io.modelcontextprotocol.spec.McpSchema.ReadResourceResult history(io.modelcontextprotocol.server.McpSyncServerExchange exchange, io.modelcontextprotocol.common.McpTransportContext transportContext, io.modelcontextprotocol.spec.McpSchema.ReadResourceRequest req) {
      var ctx = exchange != null ? (io.jooby.Context) exchange.transportContext().get("CTX") : (transportContext != null ? (io.jooby.Context) transportContext.get("CTX") : null);
      var uri = req.uri();
      var manager = new io.modelcontextprotocol.util.DefaultMcpUriTemplateManager("calculator://history/{user}");
      var args = new java.util.HashMap<String, Object>();
      args.putAll(manager.extractVariableValues(uri));
      var c = this.factory.apply(ctx);
      var raw_user = args.get("user");
      var user = raw_user != null ? raw_user.toString() : null;
      var result = c.history(user);
      return new io.jooby.mcp.McpResult(this.json).toResourceResult(req.uri(), result);
    }

    private io.modelcontextprotocol.spec.McpSchema.Tool getSessionInfoToolSpec(com.github.victools.jsonschema.generator.SchemaGenerator schemaGenerator) {
      var schema = new java.util.LinkedHashMap<String, Object>();
      schema.put("type", "object");
      var props = new java.util.LinkedHashMap<String, Object>();
      schema.put("properties", props);
      var req = new java.util.ArrayList<String>();
      schema.put("required", req);
      return new io.modelcontextprotocol.spec.McpSchema.Tool("get_session_info", null, "Returns the current MCP session ID using the injected exchange.", this.json.convertValue(schema, io.modelcontextprotocol.spec.McpSchema.JsonSchema.class), null, null, null);
    }

    private io.modelcontextprotocol.spec.McpSchema.CallToolResult getSessionInfo(io.modelcontextprotocol.server.McpSyncServerExchange exchange, io.modelcontextprotocol.common.McpTransportContext transportContext, io.modelcontextprotocol.spec.McpSchema.CallToolRequest req) {
      var ctx = exchange != null ? (io.jooby.Context) exchange.transportContext().get("CTX") : (transportContext != null ? (io.jooby.Context) transportContext.get("CTX") : null);
      var args = req.arguments() != null ? req.arguments() : java.util.Collections.<String, Object>emptyMap();
      var c = this.factory.apply(ctx);
      var result = c.getSessionInfo(exchange);
      return new io.jooby.mcp.McpResult(this.json).toCallToolResult(result, false);
    }

    private io.modelcontextprotocol.spec.McpSchema.CompleteResult historyCompletionHandler(io.modelcontextprotocol.server.McpSyncServerExchange exchange, io.modelcontextprotocol.common.McpTransportContext transportContext, io.modelcontextprotocol.spec.McpSchema.CompleteRequest req) {
      var ctx = exchange != null ? (io.jooby.Context) exchange.transportContext().get("CTX") : (transportContext != null ? (io.jooby.Context) transportContext.get("CTX") : null);
      var c = this.factory.apply(ctx);
      var targetArg = req.argument() != null ? req.argument().name() : "";
      var typedValue = req.argument() != null ? req.argument().value() : "";
      return switch (targetArg) {
        case "user" -> {
          var result = c.historyCompletions(typedValue);
          yield new io.jooby.mcp.McpResult(this.json).toCompleteResult(result);
        }
        default -> new io.jooby.mcp.McpResult(this.json).toCompleteResult(java.util.List.of());
      };
    }
}