import { createHash, randomBytes } from "node:crypto";
import { hostname } from "node:os";
import path from "node:path";

export type ExtensionInfo = {
  name: string;
  description: string;
};

type AttributeValue = string | number | boolean | Record<string, unknown>;

type UsageStats = {
  input?: number;
  output?: number;
  cacheRead?: number;
  cacheWrite?: number;
  totalTokens?: number;
  reasoning?: number;
  cost?: {
    input?: number;
    output?: number;
    cacheRead?: number;
    cacheWrite?: number;
    total?: number;
  };
};

type ModelIdentity = {
  provider?: string;
  id?: string;
};

type PiContext = {
  cwd?: string;
  model?: ModelIdentity;
  sessionManager?: {
    getSessionFile?(): string | undefined;
    getSessionId?(): string;
    getCwd?(): string;
  };
};

type PiApi = {
  on(
    eventName: string,
    handler: (
      event: Record<string, unknown>,
      ctx: PiContext,
    ) => Promise<void> | void,
  ): void;
};

type ToolCallState = {
  spanId: string;
  name: string;
  startTimeUnixNano: string;
  input: unknown;
};

export type Span = {
  traceId: string;
  spanId: string;
  parentSpanId?: string;
  name: string;
  kind: number;
  startTimeUnixNano: string;
  endTimeUnixNano: string;
  attributes: Array<{
    key: string;
    value: {
      stringValue?: string;
      intValue?: string;
      doubleValue?: number;
      boolValue?: boolean;
    };
  }>;
  status?: {
    code: number;
    message?: string;
  };
};

type OtlpAttribute = Span["attributes"][number];

type SessionInfo = {
  id: string;
  cwd: string;
  sourcePath?: string;
  host: string;
};

type ExtensionOptions = {
  endpoint?: string;
  environment?: string;
  hostName?: string;
  exporter?: (spans: Span[]) => Promise<void> | void;
};

const PACKAGE_VERSION = "0.1.1";
const DEFAULT_ENDPOINT = "https://langfuse.ai.roxasroot.net/otel/v1/traces";
const STATUS_OK = 1;
const STATUS_ERROR = 2;
const SPAN_KIND_INTERNAL = 1;

export const extensionInfo: ExtensionInfo = {
  name: "langfuse",
  description:
    "Pi extension that exports full session traces, tool activity, usage, and costs to Langfuse",
};

function nowUnixNano(): string {
  return (BigInt(Date.now()) * 1_000_000n).toString();
}

function newTraceId(): string {
  return randomBytes(16).toString("hex");
}

function newSpanId(): string {
  return randomBytes(8).toString("hex");
}

function stableSpanId(value: string): string {
  return createHash("sha256").update(value).digest("hex").slice(0, 16);
}

function asRecord(value: unknown): Record<string, unknown> | undefined {
  return typeof value === "object" && value !== null && !Array.isArray(value)
    ? (value as Record<string, unknown>)
    : undefined;
}

function asNumber(value: unknown): number | undefined {
  return typeof value === "number" && Number.isFinite(value)
    ? value
    : undefined;
}

function normalizeUsage(value: unknown): UsageStats | undefined {
  const usage = asRecord(value);
  if (!usage) return undefined;

  const cost = asRecord(usage.cost);
  const totalTokens =
    asNumber(usage.totalTokens) ?? asNumber(usage.total) ?? undefined;

  return {
    input: asNumber(usage.input),
    output: asNumber(usage.output),
    cacheRead: asNumber(usage.cacheRead),
    cacheWrite: asNumber(usage.cacheWrite),
    totalTokens,
    reasoning: asNumber(usage.reasoning),
    cost: cost
      ? {
          input: asNumber(cost.input),
          output: asNumber(cost.output),
          cacheRead: asNumber(cost.cacheRead),
          cacheWrite: asNumber(cost.cacheWrite),
          total: asNumber(cost.total),
        }
      : undefined,
  };
}

function totalTokens(usage: UsageStats | undefined): number | undefined {
  if (!usage) return undefined;
  const total =
    usage.totalTokens ??
    (usage.input ?? 0) +
      (usage.output ?? 0) +
      (usage.cacheRead ?? 0) +
      (usage.cacheWrite ?? 0);
  return total > 0 ? total : undefined;
}

function usageDetails(
  usage: UsageStats | undefined,
): Record<string, number> | undefined {
  if (!usage) return undefined;
  const details: Record<string, number> = {};
  const cacheRead = usage.cacheRead ?? 0;
  const cacheWrite = usage.cacheWrite ?? 0;
  const regularInput =
    usage.input === undefined
      ? undefined
      : Math.max(usage.input - cacheRead - cacheWrite, 0);

  if (regularInput !== undefined) details.input = regularInput;
  if (usage.output !== undefined) details.output = usage.output;
  if (usage.reasoning !== undefined) details.output_reasoning = usage.reasoning;
  if (usage.cacheRead !== undefined)
    details.input_cached_tokens = usage.cacheRead;
  if (usage.cacheWrite !== undefined)
    details.input_cache_creation = usage.cacheWrite;
  if (totalTokens(usage) !== undefined) details.total = totalTokens(usage) ?? 0;
  return Object.keys(details).length > 0 ? details : undefined;
}

function costDetails(
  usage: UsageStats | undefined,
): Record<string, number> | undefined {
  if (!usage?.cost?.total || usage.cost.total <= 0) return undefined;
  const details: Record<string, number> = { total: usage.cost.total };
  if (usage.cost.input !== undefined) details.input = usage.cost.input;
  if (usage.cost.output !== undefined) details.output = usage.cost.output;
  if (usage.cost.cacheRead !== undefined)
    details.cache_read = usage.cost.cacheRead;
  if (usage.cost.cacheWrite !== undefined)
    details.cache_write = usage.cost.cacheWrite;
  return details;
}

function attr(
  key: string,
  value: AttributeValue | undefined,
): OtlpAttribute | undefined {
  if (value === undefined || value === "") return undefined;
  if (typeof value === "boolean") return { key, value: { boolValue: value } };
  if (typeof value === "number") {
    if (!Number.isFinite(value)) return undefined;
    if (Number.isInteger(value))
      return { key, value: { intValue: value.toString() } };
    return { key, value: { doubleValue: value } };
  }
  if (typeof value === "object") {
    return { key, value: { stringValue: stringifySample(value, 8000) ?? "" } };
  }
  return { key, value: { stringValue: value } };
}

function attrs(
  values: Record<string, AttributeValue | undefined>,
): Span["attributes"] {
  return Object.entries(values)
    .map(([key, value]) => attr(key, value))
    .filter((value): value is OtlpAttribute => value !== undefined);
}

function stringifySample(
  value: unknown,
  maxLength = 16_000,
): string | undefined {
  if (value === undefined) return undefined;
  let text: string;
  try {
    text = typeof value === "string" ? value : JSON.stringify(value);
  } catch {
    text = String(value);
  }
  return text.length > maxLength
    ? `${text.slice(0, maxLength)}\n[truncated by pi-langfuse]`
    : text;
}

function contentToText(value: unknown): string | undefined {
  if (typeof value === "string") return value;
  if (!Array.isArray(value)) return stringifySample(value);
  const parts = value
    .map((item) => {
      const part = asRecord(item);
      if (part?.type === "text" && typeof part.text === "string") {
        return part.text;
      }
      if (part?.type === "thinking" && typeof part.thinking === "string") {
        return part.thinking;
      }
      if (part?.type === "toolCall") return stringifySample(part);
      if (part?.type === "image") return "[image]";
      return stringifySample(part);
    })
    .filter((item): item is string => Boolean(item));
  return parts.length > 0 ? parts.join("\n") : undefined;
}

function messageContent(value: Record<string, unknown>): unknown {
  return value.content;
}

function messageTimestamp(value: Record<string, unknown>): string {
  const timestamp = asNumber(value.timestamp);
  return timestamp === undefined
    ? nowUnixNano()
    : `${BigInt(timestamp) * 1_000_000n}`;
}

function normalizeModelName(model: string | undefined): string | undefined {
  if (!model) return undefined;
  if (model === "gpt-5.4-mini") return undefined;
  if (model.startsWith("kimi-for-coding")) return "kimi-for-coding";
  return model;
}

function projectMetadata(cwd: string): {
  projectPath: string;
  projectName?: string;
  projectFolder?: string;
} {
  const normalized = cwd.replace(/[\\/]+$/, "");
  const projectName = normalized
    .split(/[\\/]+/)
    .filter(Boolean)
    .at(-1);
  return {
    projectPath: cwd,
    projectName,
    projectFolder: projectName,
  };
}

function sessionInfo(ctx: PiContext, hostName: string): SessionInfo {
  const sourcePath = ctx.sessionManager?.getSessionFile?.();
  const id =
    ctx.sessionManager?.getSessionId?.() ??
    (sourcePath
      ? path.basename(sourcePath, path.extname(sourcePath))
      : "pi-live");
  const cwd = ctx.sessionManager?.getCwd?.() ?? ctx.cwd ?? process.cwd();

  return {
    id,
    cwd,
    sourcePath,
    host: hostName,
  };
}

function commonAttributes(
  info: SessionInfo,
  recordId: string,
  environment: string,
  model?: string,
  provider?: string,
): Record<string, AttributeValue | undefined> {
  const project = projectMetadata(info.cwd);
  return {
    "service.name": "agent.pi",
    "deployment.environment": environment,
    "langfuse.environment": environment,
    "langfuse.session.id": info.id,
    "session.id": info.id,
    "agent.name": "pi",
    "host.name": info.host,
    "agent.session_id": info.id,
    "agent.record_id": recordId,
    "source.path": info.sourcePath,
    cwd: info.cwd,
    "project.path": project.projectPath,
    "project.name": project.projectName,
    "project.folder": project.projectFolder,
    "telemetry.source": "pi-langfuse-extension",
    "langfuse.observation.metadata.agent": "pi",
    "langfuse.observation.metadata.host": info.host,
    "langfuse.observation.metadata.machine": info.host,
    "langfuse.observation.metadata.session_id": info.id,
    "langfuse.observation.metadata.record_id": recordId,
    "langfuse.observation.metadata.source_path": info.sourcePath,
    "langfuse.observation.metadata.cwd": info.cwd,
    "langfuse.observation.metadata.project_path": project.projectPath,
    "langfuse.observation.metadata.project_name": project.projectName,
    "langfuse.observation.metadata.project_folder": project.projectFolder,
    "langfuse.observation.metadata.model": model,
    "langfuse.observation.metadata.provider": provider,
    "langfuse.observation.metadata.telemetry_source": "pi-langfuse-extension",
    "agent.model": model,
    "agent.provider": provider,
  };
}

function rootSpan(input: {
  traceId: string;
  rootSpanId: string;
  startTimeUnixNano: string;
  info: SessionInfo;
  environment: string;
}): Span {
  const project = projectMetadata(input.info.cwd);
  return makeSpan({
    traceId: input.traceId,
    spanId: input.rootSpanId,
    name: "pi session",
    startTimeUnixNano: input.startTimeUnixNano,
    endTimeUnixNano: input.startTimeUnixNano,
    attributes: {
      ...commonAttributes(input.info, "session-root", input.environment),
      "langfuse.internal.as_root": true,
      "langfuse.trace.name": "pi session",
      "langfuse.trace.metadata.agent": "pi",
      "langfuse.trace.metadata.host": input.info.host,
      "langfuse.trace.metadata.machine": input.info.host,
      "langfuse.trace.metadata.source_path": input.info.sourcePath,
      "langfuse.trace.metadata.cwd": input.info.cwd,
      "langfuse.trace.metadata.project_path": project.projectPath,
      "langfuse.trace.metadata.project_name": project.projectName,
      "langfuse.trace.metadata.project_folder": project.projectFolder,
      "langfuse.trace.metadata.telemetry_source": "pi-langfuse-extension",
      "langfuse.observation.type": "span",
    },
  });
}

function makeSpan(input: {
  traceId: string;
  spanId?: string;
  parentSpanId?: string;
  name: string;
  startTimeUnixNano?: string;
  endTimeUnixNano?: string;
  attributes?: Record<string, AttributeValue | undefined>;
  statusCode?: number;
  statusMessage?: string;
}): Span {
  return {
    traceId: input.traceId,
    spanId: input.spanId ?? newSpanId(),
    parentSpanId: input.parentSpanId,
    name: input.name,
    kind: SPAN_KIND_INTERNAL,
    startTimeUnixNano: input.startTimeUnixNano ?? nowUnixNano(),
    endTimeUnixNano: input.endTimeUnixNano ?? nowUnixNano(),
    attributes: attrs(input.attributes ?? {}),
    status: {
      code: input.statusCode ?? STATUS_OK,
      message: input.statusMessage,
    },
  };
}

export function buildOtlpTracePayload(spans: Span[], hostName = hostname()) {
  return {
    resourceSpans: [
      {
        resource: {
          attributes: attrs({
            "service.name": "agent.pi",
            "service.namespace": "toolbox",
            "host.name": hostName,
            "telemetry.source": "pi-langfuse-extension",
          }),
        },
        scopeSpans: [
          {
            scope: {
              name: "@ramarivera/pi-langfuse",
              version: PACKAGE_VERSION,
            },
            spans,
          },
        ],
      },
    ],
  };
}

async function exportSpans(
  spans: Span[],
  options: ExtensionOptions,
): Promise<void> {
  if (process.env.PI_LANGFUSE_ENABLED === "false" || spans.length === 0) return;
  if (options.exporter) {
    await options.exporter(spans);
    return;
  }

  const endpoint =
    options.endpoint ??
    process.env.PI_LANGFUSE_OTLP_ENDPOINT ??
    process.env.OTEL_EXPORTER_OTLP_TRACES_ENDPOINT ??
    DEFAULT_ENDPOINT;
  try {
    await fetch(endpoint, {
      method: "POST",
      headers: { "content-type": "application/json" },
      body: JSON.stringify(buildOtlpTracePayload(spans, options.hostName)),
    });
  } catch (error) {
    if (process.env.PI_LANGFUSE_DEBUG === "true") {
      console.error("[pi-langfuse] OTLP export failed", error);
    }
  }
}

export function createExtension(options: ExtensionOptions = {}) {
  return function piLangfuse(pi: PiApi) {
    const environment =
      options.environment ?? process.env.PI_LANGFUSE_ENVIRONMENT ?? "local";
    const hostName = options.hostName ?? hostname();
    let traceId = newTraceId();
    let rootSpanId = stableSpanId(`${traceId}:root`);
    let rootStartTimeUnixNano = nowUnixNano();
    let rootEmitted = false;
    const emittedRecords = new Set<string>();
    const toolCalls = new Map<string, ToolCallState>();

    const emit = async (spans: Span[]) =>
      exportSpans(spans, { ...options, hostName });

    const ensureRoot = async (ctx: PiContext) => {
      if (rootEmitted) return;
      rootEmitted = true;
      await emit([
        rootSpan({
          traceId,
          rootSpanId,
          startTimeUnixNano: rootStartTimeUnixNano,
          info: sessionInfo(ctx, hostName),
          environment,
        }),
      ]);
    };

    pi.on("session_start", async (_event, ctx) => {
      traceId = newTraceId();
      rootSpanId = stableSpanId(`${traceId}:root`);
      rootStartTimeUnixNano = nowUnixNano();
      rootEmitted = false;
      emittedRecords.clear();
      toolCalls.clear();
      await ensureRoot(ctx);
    });

    pi.on("agent_start", async (_event, ctx) => {
      await ensureRoot(ctx);
    });

    pi.on("message_end", async (event, ctx) => {
      await ensureRoot(ctx);
      const message = asRecord(event.message);
      if (!message) return;

      const role = typeof message.role === "string" ? message.role : "unknown";
      const recordId = `${role}:${message.timestamp ?? ""}:${createHash(
        "sha256",
      )
        .update(stringifySample(messageContent(message)) ?? "")
        .digest("hex")
        .slice(0, 16)}`;
      if (emittedRecords.has(recordId)) return;
      emittedRecords.add(recordId);

      const info = sessionInfo(ctx, hostName);
      const usage = normalizeUsage(message.usage);
      const provider =
        typeof message.provider === "string"
          ? message.provider
          : ctx.model?.provider;
      const model = normalizeModelName(
        typeof message.model === "string" ? message.model : ctx.model?.id,
      );
      const generation = role === "assistant";
      const content = messageContent(message);

      await emit([
        makeSpan({
          traceId,
          spanId: stableSpanId(`${traceId}:${recordId}`),
          parentSpanId: rootSpanId,
          name: generation ? "pi assistant" : `pi ${role}`,
          startTimeUnixNano: messageTimestamp(message),
          endTimeUnixNano: nowUnixNano(),
          attributes: {
            ...commonAttributes(info, recordId, environment, model, provider),
            "langfuse.observation.type": generation ? "generation" : "span",
            "langfuse.observation.model.name": generation ? model : undefined,
            "langfuse.observation.usage_details": usageDetails(usage),
            "langfuse.observation.cost_details": costDetails(usage),
            "langfuse.observation.input":
              role === "user" ? contentToText(content) : undefined,
            "langfuse.observation.output":
              role === "user" ? undefined : contentToText(content),
            "gen_ai.response.model": generation ? model : undefined,
            "gen_ai.usage.input_tokens": usage?.input,
            "gen_ai.usage.output_tokens": usage?.output,
            "gen_ai.usage.total_tokens": totalTokens(usage),
            "gen_ai.usage.cache_read_tokens": usage?.cacheRead,
            "gen_ai.usage.cache_write_tokens": usage?.cacheWrite,
            "gen_ai.usage.cost": usage?.cost?.total,
            role,
            "input.value": role === "user" ? contentToText(content) : undefined,
            "output.value":
              role === "user" ? undefined : contentToText(content),
            metadata: {
              role,
              stopReason: message.stopReason,
              responseId: message.responseId,
              content,
            },
          },
          statusCode: message.errorMessage ? STATUS_ERROR : STATUS_OK,
          statusMessage:
            typeof message.errorMessage === "string"
              ? message.errorMessage
              : undefined,
        }),
      ]);
    });

    pi.on("tool_call", async (event, ctx) => {
      await ensureRoot(ctx);
      const toolCallId =
        typeof event.toolCallId === "string" ? event.toolCallId : newSpanId();
      const toolName =
        typeof event.toolName === "string" ? event.toolName : "tool";
      toolCalls.set(toolCallId, {
        spanId: stableSpanId(`${traceId}:tool:${toolCallId}`),
        name: toolName,
        startTimeUnixNano: nowUnixNano(),
        input: event.input,
      });
    });

    pi.on("tool_result", async (event, ctx) => {
      await ensureRoot(ctx);
      const info = sessionInfo(ctx, hostName);
      const toolCallId =
        typeof event.toolCallId === "string" ? event.toolCallId : newSpanId();
      const state = toolCalls.get(toolCallId);
      const toolName =
        state?.name ??
        (typeof event.toolName === "string" ? event.toolName : "tool");
      const recordId = `tool:${toolCallId}`;
      if (emittedRecords.has(recordId)) return;
      emittedRecords.add(recordId);
      const isError = event.isError === true;

      await emit([
        makeSpan({
          traceId,
          spanId:
            state?.spanId ?? stableSpanId(`${traceId}:tool:${toolCallId}`),
          parentSpanId: rootSpanId,
          name: `pi tool ${toolName}`,
          startTimeUnixNano: state?.startTimeUnixNano,
          attributes: {
            ...commonAttributes(info, recordId, environment),
            "langfuse.observation.type": "span",
            "langfuse.observation.input": stringifySample(
              state?.input ?? event.input,
            ),
            "langfuse.observation.output": contentToText(event.content),
            "gen_ai.system": "pi",
            "gen_ai.tool.name": toolName,
            "tool.call.id": toolCallId,
            "tool.input": stringifySample(state?.input ?? event.input),
            "tool.result": contentToText(event.content),
            "tool.success": !isError,
            role: "tool",
            metadata: {
              toolName,
              toolCallId,
              details: event.details,
              isError,
            },
          },
          statusCode: isError ? STATUS_ERROR : STATUS_OK,
        }),
      ]);

      toolCalls.delete(toolCallId);
    });

    pi.on("agent_end", async (_event, ctx) => {
      await ensureRoot(ctx);
    });
  };
}