import type { Static, TSchema } from "typebox";
import { Check, Convert } from "typebox/value";
import { WorkflowError, WorkflowErrorCode } from "./errors.js";

export interface StructuredOutputCapture<T = unknown> {
  value: T | undefined;
  called: boolean;
}

export interface StructuredOutputToolOptions<TSchemaDef extends TSchema> {
  schema: TSchemaDef;
  capture: StructuredOutputCapture<Static<TSchemaDef>>;
  name?: string;
}

/**
 * Structural tool shape understood by Pi's custom-tool API.
 *
 * Keeping this tiny contract local is intentional: `defineTool()` is an identity
 * helper, and importing it here made the Claude/Codex-only CLI load the entire Pi
 * SDK before it could even print `--help`.
 */
export interface StructuredOutputTool<TSchemaDef extends TSchema> {
  name: string;
  label: string;
  description: string;
  promptSnippet: string;
  promptGuidelines: string[];
  parameters: TSchemaDef;
  execute(
    toolCallId: string,
    params: Static<TSchemaDef>,
  ): Promise<{
    content: Array<{ type: "text"; text: string }>;
    details: Static<TSchemaDef>;
    terminate: true;
  }>;
}

/**
 * Create a terminating tool that captures validated params as the subagent result.
 *
 * Pi validates `params` against `schema` before execute() is called. Returning
 * `terminate: true` lets the subagent finish on this tool call without paying for
 * an extra assistant follow-up turn.
 */
export function createStructuredOutputTool<TSchemaDef extends TSchema>({
  schema,
  capture,
  name = "structured_output",
}: StructuredOutputToolOptions<TSchemaDef>): StructuredOutputTool<TSchemaDef> {
  return {
    name,
    label: "Structured Output",
    description: "Return the final machine-readable result for this subagent task.",
    promptSnippet: "Return final machine-readable output",
    promptGuidelines: [
      `${name} is the final answer channel for this task; call ${name} exactly once when done.`,
      `Do not write a prose final answer after calling ${name}.`,
    ],
    parameters: schema,
    async execute(_toolCallId, params) {
      capture.value = params;
      capture.called = true;
      return {
        content: [{ type: "text", text: "Structured output received." }],
        details: params,
        terminate: true,
      };
    },
  };
}

/** A schema accepted by the CLI providers (JSON Schema/TypeBox compatible). */
export type StructuredOutputSchema = TSchema & { type?: string };

/** Ensure provider structured-output schemas are transport-safe top-level objects. */
export function assertTopLevelObjectSchema(
  schema: unknown,
  context = "structured output",
): asserts schema is StructuredOutputSchema {
  if (!schema || typeof schema !== "object" || (schema as { type?: unknown }).type !== "object") {
    const type = schema && typeof schema === "object" ? (schema as { type?: unknown }).type : undefined;
    throw new WorkflowError(
      `${context} schema must be a top-level JSON object schema (type: "object") — got type: ${type ?? "undefined"}`,
      WorkflowErrorCode.SCHEMA_NONCOMPLIANCE,
      { recoverable: false, details: { schema } },
    );
  }
}

/** Validate and TypeBox-convert a value, throwing a non-recoverable schema error. */
export function validateStructuredOutput<T = unknown>(
  value: unknown,
  schema: TSchema,
  context = "structured output",
): T {
  assertTopLevelObjectSchema(schema, context);
  try {
    const converted = Convert(schema, value);
    if (Check(schema, converted)) return converted as T;
  } catch (error) {
    throw new WorkflowError(`${context} did not satisfy its schema`, WorkflowErrorCode.SCHEMA_NONCOMPLIANCE, {
      recoverable: false,
      details: { value, schema, cause: error },
    });
  }
  throw new WorkflowError(`${context} did not satisfy its schema`, WorkflowErrorCode.SCHEMA_NONCOMPLIANCE, {
    recoverable: false,
    details: { value, schema },
  });
}

/** Parse exactly one JSON value and validate it against a top-level object schema. */
export function parseStructuredOutput<T = unknown>(text: string, schema: TSchema, context = "structured output"): T {
  assertTopLevelObjectSchema(schema, context);
  if (typeof text !== "string" || !text.trim()) {
    throw new WorkflowError(`${context} was empty`, WorkflowErrorCode.SCHEMA_NONCOMPLIANCE, {
      recoverable: false,
      details: { text },
    });
  }
  let value: unknown;
  try {
    value = JSON.parse(text);
  } catch (error) {
    throw new WorkflowError(`${context} was not valid JSON`, WorkflowErrorCode.SCHEMA_NONCOMPLIANCE, {
      recoverable: false,
      details: { text, cause: error },
    });
  }
  return validateStructuredOutput<T>(value, schema, context);
}

/**
 * Locate a JSON object in a model response. Fenced JSON is preferred, then the
 * first balanced object is considered. Strings and escapes are handled so braces
 * in prose or string values do not terminate extraction early.
 */
export function extractJsonObject(text: string): string | undefined {
  const fenced = text.match(/```(?:json)?\s*([\s\S]*?)```/i)?.[1]?.trim();
  if (fenced) return fenced;
  for (let start = 0; start < text.length; start++) {
    if (text[start] !== "{") continue;
    let depth = 0;
    let inString = false;
    let escaped = false;
    for (let index = start; index < text.length; index++) {
      const char = text[index];
      if (inString) {
        if (escaped) escaped = false;
        else if (char === "\\") escaped = true;
        else if (char === '"') inString = false;
        continue;
      }
      if (char === '"') {
        inString = true;
        continue;
      }
      if (char === "{") depth++;
      else if (char === "}" && --depth === 0) return text.slice(start, index + 1);
    }
  }
  return undefined;
}

/** Strictly extract and validate an object embedded in model prose. */
export function extractStructuredOutput<T = unknown>(text: string, schema: TSchema, context = "structured output"): T {
  assertTopLevelObjectSchema(schema, context);
  const json = extractJsonObject(text);
  if (!json) {
    throw new WorkflowError(
      `${context} could not be extracted from model text`,
      WorkflowErrorCode.SCHEMA_NONCOMPLIANCE,
      {
        recoverable: false,
        details: { text },
      },
    );
  }
  return parseStructuredOutput<T>(json, schema, context);
}

/** Parse a provider result that may already be an object or may be JSON text. */
export function resolveStructuredOutput<T = unknown>(
  value: unknown,
  schema: TSchema,
  context = "structured output",
): T {
  assertTopLevelObjectSchema(schema, context);
  if (typeof value === "string") {
    try {
      return parseStructuredOutput<T>(value, schema, context);
    } catch (strictError) {
      try {
        return extractStructuredOutput<T>(value, schema, context);
      } catch {
        throw strictError;
      }
    }
  }
  return validateStructuredOutput<T>(value, schema, context);
}