import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
import { Type, type Static } from "typebox";
import { execFile } from "node:child_process";
import { mkdtemp, rm, writeFile } from "node:fs/promises";
import { existsSync } from "node:fs";
import { join, resolve } from "node:path";
import { tmpdir } from "node:os";

const MODEL_CANDIDATES = [
  "semantic.yml",
  "semantic.yaml",
  "model.yml",
  "model.yaml",
  "semantic_layer.yml",
  "semantic_layer.yaml",
];

const BaseParams = Type.Object({
  model: Type.Optional(
    Type.String({ description: "Path to semantic model YAML. Defaults to semantic.yml/model.yml candidates in the current project." }),
  ),
  semanticBinary: Type.Optional(
    Type.String({ description: "Semantic CLI binary to run. Defaults to semantic." }),
  ),
});

const MetricsParams = BaseParams;

const SearchParams = Type.Intersect([
  BaseParams,
  Type.Object({
    query: Type.String({ description: "Metric search query. Matches id, name, description, synonyms, and teams." }),
    limit: Type.Optional(Type.Number({ description: "Maximum number of matches. Defaults to 10." })),
  }),
]);

const DescribeParams = Type.Intersect([
  BaseParams,
  Type.Object({
    metricId: Type.String({ description: "Metric identifier to describe." }),
  }),
]);

const ValidateParams = Type.Intersect([
  BaseParams,
  Type.Object({
    dialect: Type.Optional(Type.String({ description: "Dialect for compile checks. Defaults to bigquery." })),
    checkCompilable: Type.Optional(Type.Boolean({ description: "Compile every metric as part of validation." })),
    columnRegistry: Type.Optional(Type.String({ description: "Optional column registry JSON path." })),
    maxRegistryAgeHours: Type.Optional(Type.Number({ description: "Warn when registry metadata is older than this many hours." })),
    requireJoinUniqueness: Type.Optional(Type.Boolean({ description: "Require join target uniqueness metadata from the registry." })),
  }),
]);

const CompileParams = Type.Intersect([
  BaseParams,
  Type.Object({
    metricId: Type.String({ description: "Metric identifier to compile." }),
    dialect: Type.Optional(Type.String({ description: "Target SQL dialect. Defaults to bigquery." })),
    period: Type.Optional(Type.String({ description: "Named period, e.g. current year or last 12 complete months." })),
    requestPath: Type.Optional(Type.String({ description: "Path to request JSON shape." })),
    request: Type.Optional(Type.Any({ description: "Inline request JSON shape. Use this instead of requestPath for agent-created requests." })),
    format: Type.Optional(Type.Union([
      Type.Literal("sql"),
      Type.Literal("json"),
      Type.Literal("explain"),
    ], { description: "Output format. Defaults to sql." })),
    targetTable: Type.Optional(Type.String({ description: "Optional target table for targetSeries requests. May be fully qualified, or bare when targetProject/targetSchema are supplied." })),
    targetProject: Type.Optional(Type.String({ description: "Optional target table project/catalog." })),
    targetSchema: Type.Optional(Type.String({ description: "Optional target table schema/dataset." })),
    targetMetricColumn: Type.Optional(Type.String({ description: "Target metric id column. Defaults to metric_id." })),
    targetSeriesColumn: Type.Optional(Type.String({ description: "Target series column. Defaults to target_series." })),
    targetTimeColumn: Type.Optional(Type.String({ description: "Target time column. Defaults to metric_time." })),
    targetValueColumn: Type.Optional(Type.String({ description: "Target value column. Defaults to target_value." })),
    targetDimensionColumns: Type.Optional(Type.Array(Type.String(), { description: "Target dimension columns used to infer the lowest matching target grain." })),
    targetNullColumns: Type.Optional(Type.Array(Type.String(), { description: "Target columns that must be NULL for this request." })),
  }),
]);

type BaseParamsType = Static<typeof BaseParams>;
type SearchParamsType = Static<typeof SearchParams>;
type DescribeParamsType = Static<typeof DescribeParams>;
type ValidateParamsType = Static<typeof ValidateParams>;
type CompileParamsType = Static<typeof CompileParams>;

type SemanticResult = {
  stdout: string;
  stderr: string;
  command: string[];
};

function semanticBinary(params: BaseParamsType): string {
  return params.semanticBinary || "semantic";
}

function resolveModelPath(ctx: ExtensionContext, model?: string): string {
  if (model) return resolve(ctx.cwd, model);
  for (const candidate of MODEL_CANDIDATES) {
    const path = resolve(ctx.cwd, candidate);
    if (existsSync(path)) return path;
  }
  throw new Error(
    `No semantic model found. Pass model explicitly or add one of: ${MODEL_CANDIDATES.join(", ")}`,
  );
}

function runSemantic(
  ctx: ExtensionContext,
  params: BaseParamsType,
  args: string[],
  signal: AbortSignal | undefined,
): Promise<SemanticResult> {
  const command = semanticBinary(params);
  return new Promise((resolvePromise, reject) => {
    execFile(command, args, { cwd: ctx.cwd, signal, maxBuffer: 20 * 1024 * 1024 }, (error, stdout, stderr) => {
      const result = { stdout, stderr, command: [command, ...args] };
      if (error) {
        const message = stderr.trim() || error.message;
        if (message.includes("unrecognized arguments") && (message.includes("--target-project") || message.includes("--target-schema"))) {
          reject(new Error(
            `${command} failed because targetProject/targetSchema require semantic-query-compiler >= 0.1.1. `
              + `Upgrade the Python CLI with: uv tool install 'semantic-query-compiler>=0.1.4' --force --refresh. Original error: ${message}`,
          ));
          return;
        }
        reject(new Error(`${command} failed: ${message}`));
        return;
      }
      resolvePromise(result);
    });
  });
}

function jsonContent(result: SemanticResult) {
  const text = result.stdout.trim();
  let parsed: unknown;
  try {
    parsed = text ? JSON.parse(text) : null;
  } catch {
    parsed = text;
  }
  return {
    content: [{ type: "text" as const, text: text || "{}" }],
    details: { command: result.command, result: parsed, stderr: result.stderr || undefined },
  };
}

function textContent(result: SemanticResult) {
  return {
    content: [{ type: "text" as const, text: result.stdout }],
    details: { command: result.command, stderr: result.stderr || undefined },
  };
}

async function withInlineRequest<T>(
  ctx: ExtensionContext,
  request: unknown,
  callback: (path: string | undefined) => Promise<T>,
): Promise<T> {
  if (request === undefined) return callback(undefined);
  const dir = await mkdtemp(join(tmpdir(), "pi-semantic-query-"));
  const path = join(dir, "request.json");
  try {
    await writeFile(path, JSON.stringify(request, null, 2), "utf8");
    return await callback(path);
  } finally {
    await rm(dir, { recursive: true, force: true });
  }
}

export default function semanticQueryExtension(pi: ExtensionAPI) {
  pi.registerTool({
    name: "semantic_metrics",
    label: "Semantic Metrics",
    description: "List semantic metrics with discovery metadata from the current project model.",
    promptSnippet: "Use semantic_metrics to inspect available governed metrics before compiling SQL.",
    promptGuidelines: [
      "Use this before guessing metric identifiers.",
      "Prefer semantic_search_metrics when the user describes a metric in natural language.",
    ],
    parameters: MetricsParams,
    async execute(_toolCallId, params: BaseParamsType, signal, _onUpdate, ctx) {
      const model = resolveModelPath(ctx, params.model);
      return jsonContent(await runSemantic(ctx, params, ["metrics", "--model", model, "--format", "json"], signal));
    },
  });

  pi.registerTool({
    name: "semantic_search_metrics",
    label: "Semantic Search Metrics",
    description: "Search semantic metrics by id, name, description, synonyms, and teams.",
    promptSnippet: "Use semantic_search_metrics to find the right governed metric from a natural-language phrase.",
    promptGuidelines: ["Use this when a user asks for a metric by business wording rather than exact id."],
    parameters: SearchParams,
    async execute(_toolCallId, params: SearchParamsType, signal, _onUpdate, ctx) {
      const model = resolveModelPath(ctx, params.model);
      const args = [
        "search-metrics",
        params.query,
        "--model",
        model,
        "--format",
        "json",
      ];
      if (params.limit !== undefined) args.push("--limit", String(params.limit));
      return jsonContent(await runSemantic(ctx, params, args, signal));
    },
  });

  pi.registerTool({
    name: "semantic_describe",
    label: "Semantic Describe",
    description: "Describe one semantic metric, including filters, dimensions, synonyms, and formula notes.",
    promptSnippet: "Use semantic_describe before compiling a metric if the grain, filters, or semantics are unclear.",
    parameters: DescribeParams,
    async execute(_toolCallId, params: DescribeParamsType, signal, _onUpdate, ctx) {
      const model = resolveModelPath(ctx, params.model);
      return jsonContent(await runSemantic(ctx, params, ["describe", params.metricId, "--model", model, "--format", "json"], signal));
    },
  });

  pi.registerTool({
    name: "semantic_validate",
    label: "Semantic Validate",
    description: "Validate a semantic model and optionally compile-check every metric.",
    promptSnippet: "Use semantic_validate before relying on model changes or generated SQL.",
    parameters: ValidateParams,
    async execute(_toolCallId, params: ValidateParamsType, signal, _onUpdate, ctx) {
      const model = resolveModelPath(ctx, params.model);
      const args = ["validate", "--model", model, "--format", "json"];
      if (params.dialect) args.push("--dialect", params.dialect);
      if (params.checkCompilable) args.push("--check-compilable");
      if (params.columnRegistry) args.push("--column-registry", resolve(ctx.cwd, params.columnRegistry));
      if (params.maxRegistryAgeHours !== undefined) args.push("--max-registry-age-hours", String(params.maxRegistryAgeHours));
      if (params.requireJoinUniqueness) args.push("--require-join-uniqueness");
      return jsonContent(await runSemantic(ctx, params, args, signal));
    },
  });

  pi.registerTool({
    name: "semantic_compile",
    label: "Semantic Compile",
    description: "Compile a governed semantic metric request to SQL. Does not execute warehouse queries.",
    promptSnippet: "Use semantic_compile to produce inspectable SQL for a governed metric. It does not execute the query.",
    promptGuidelines: [
      "Prefer period over hard-coded fromDate/toDate for CLI use.",
      "Use inline request for agent-created request shapes.",
      "Do not present compiled SQL as executed results.",
    ],
    parameters: CompileParams,
    async execute(_toolCallId, params: CompileParamsType, signal, _onUpdate, ctx) {
      const model = resolveModelPath(ctx, params.model);
      return withInlineRequest(ctx, params.request, async (inlineRequestPath) => {
        const requestPath = params.requestPath ? resolve(ctx.cwd, params.requestPath) : inlineRequestPath;
        const args = ["compile", params.metricId, "--model", model];
        if (requestPath) args.push("--request", requestPath);
        if (params.period) args.push("--period", params.period);
        if (params.dialect) args.push("--dialect", params.dialect);
        if (params.format) args.push("--format", params.format);
        if (params.targetTable) args.push("--target-table", params.targetTable);
        if (params.targetProject) args.push("--target-project", params.targetProject);
        if (params.targetSchema) args.push("--target-schema", params.targetSchema);
        if (params.targetMetricColumn) args.push("--target-metric-column", params.targetMetricColumn);
        if (params.targetSeriesColumn) args.push("--target-series-column", params.targetSeriesColumn);
        if (params.targetTimeColumn) args.push("--target-time-column", params.targetTimeColumn);
        if (params.targetValueColumn) args.push("--target-value-column", params.targetValueColumn);
        for (const column of params.targetDimensionColumns || []) args.push("--target-dimension-column", column);
        for (const column of params.targetNullColumns || []) args.push("--target-null-column", column);
        const result = await runSemantic(ctx, params, args, signal);
        return params.format === "json" || params.format === "explain" ? jsonContent(result) : textContent(result);
      });
    },
  });
}