import { convertToLlm, serializeConversation } from "@mariozechner/pi-coding-agent";
import { completeSimple } from "@mariozechner/pi-ai";
import {
  needsDeveloperToSystemPatch,
  patchDeveloperToSystem,
  patchOpenAIResponsesPayload,
  shouldPatchProxyOpenAIResponsesModel,
} from "../../src/shared/openai-responses";

const SUMMARIZATION_SYSTEM_PROMPT = `You are a context summarization assistant. Your task is to read a conversation between a user and an AI coding assistant, then produce a structured summary following the exact format specified.

Do NOT continue the conversation. Do NOT respond to any questions in the conversation. ONLY output the structured summary.`;

const SUMMARIZATION_PROMPT = `The messages above are a conversation to summarize. Create a structured context checkpoint summary that another LLM will use to continue the work.

Some messages may be synthetic compression artifacts inserted by Vera before summarization, for example:
- [History snip summary]
- [Older tool result content cleared by microcompact]
- [Context collapse projection]

Treat these as lossy aids about earlier context. Preserve the factual signals they contain, but do NOT mistake them for new user requests or fresh assistant commitments.

Use this EXACT format:

## Goal
[What is the user trying to accomplish? Can be multiple items if the session covers different tasks.]

## Constraints & Preferences
- [Any constraints, preferences, or requirements mentioned by user]
- [Or "(none)" if none were mentioned]

## Progress
### Done
- [x] [Completed tasks/changes]

### In Progress
- [ ] [Current work]

### Blocked
- [Issues preventing progress, if any]

## Key Decisions
- **[Decision]**: [Brief rationale]

## Next Steps
1. [Ordered list of what should happen next]

## Critical Context
- [Any data, examples, or references needed to continue]
- [Or "(none)" if not applicable]

Keep each section concise. Preserve exact file paths, function names, and error messages.`;

const UPDATE_SUMMARIZATION_PROMPT = `The messages above are NEW conversation messages to incorporate into the existing summary provided in <previous-summary> tags.

Some messages may be synthetic compression artifacts inserted by Vera before summarization, for example:
- [History snip summary]
- [Older tool result content cleared by microcompact]
- [Context collapse projection]

Treat these as lossy aids about earlier context. Preserve the factual signals they contain, but do NOT mistake them for new user requests or fresh assistant commitments.

Update the existing structured summary with new information. RULES:
- PRESERVE all existing information from the previous summary
- ADD new progress, decisions, and context from the new messages
- UPDATE the Progress section: move items from "In Progress" to "Done" when completed
- UPDATE "Next Steps" based on what was accomplished
- PRESERVE exact file paths, function names, and error messages
- If something is no longer relevant, you may remove it

Use this EXACT format:

## Goal
[Preserve existing goals, add new ones if the task expanded]

## Constraints & Preferences
- [Preserve existing, add new ones discovered]

## Progress
### Done
- [x] [Include previously done items AND newly completed items]

### In Progress
- [ ] [Current work - update based on progress]

### Blocked
- [Current blockers - remove if resolved]

## Key Decisions
- **[Decision]**: [Brief rationale] (preserve all previous, add new)

## Next Steps
1. [Update based on current state]

## Critical Context
- [Preserve important context, add new if needed]

Keep each section concise. Preserve exact file paths, function names, and error messages.`;

function summaryPrompt(customInstructions: string | undefined, previousSummary: string | undefined): string {
  const basePrompt = previousSummary ? UPDATE_SUMMARIZATION_PROMPT : SUMMARIZATION_PROMPT;
  if (!customInstructions?.trim()) return basePrompt;
  return `${basePrompt}\n\nAdditional focus: ${customInstructions.trim()}`;
}

/**
 * Strip inline thinking blocks that some models emit when reasoning_split is
 * not enabled. Handles known tags: <think>, <thinking>, <internal_monologue>.
 */
const INLINE_THINKING_RE = /<(think|thinking|internal_monologue)>[\s\S]*?<\/\1>/gi;

function stripInlineThinking(text: string): string {
  return text.replace(INLINE_THINKING_RE, "").trim();
}

function textFromAssistantMessage(message: any): string {
  const raw = (message?.content ?? [])
    .filter((block: any) => block?.type === "text")
    .map((block: any) => block.text ?? "")
    .join("\n")
    .trim();
  return stripInlineThinking(raw);
}

function buildPayloadPatcher(model: any): ((payload: any) => any) | undefined {
  const patchResponses = shouldPatchProxyOpenAIResponsesModel(model);
  const patchDeveloper = needsDeveloperToSystemPatch(model);

  if (!patchResponses && !patchDeveloper) return undefined;

  return (payload: any) => {
    if (patchResponses) return patchOpenAIResponsesPayload(payload, SUMMARIZATION_SYSTEM_PROMPT);
    if (patchDeveloper) return patchDeveloperToSystem(payload);
    return payload;
  };
}

export async function generateStructuredSummary(
  messages: any[],
  model: any,
  reserveTokens: number,
  auth: { apiKey?: string; headers?: Record<string, string> },
  signal?: AbortSignal,
  customInstructions?: string,
  previousSummary?: string,
): Promise<string> {
  const llmMessages = convertToLlm(messages);
  const conversationText = serializeConversation(llmMessages);
  let promptText = `<conversation>\n${conversationText}\n</conversation>\n\n`;
  if (previousSummary) {
    promptText += `<previous-summary>\n${previousSummary}\n</previous-summary>\n\n`;
  }
  promptText += summaryPrompt(customInstructions, previousSummary);

  const response = await completeSimple(
    model,
    {
      systemPrompt: SUMMARIZATION_SYSTEM_PROMPT,
      messages: [
        {
          role: "user",
          content: [{ type: "text", text: promptText }],
          timestamp: Date.now(),
        },
      ],
    },
    {
      maxTokens: Math.floor(0.8 * reserveTokens),
      signal,
      apiKey: auth.apiKey,
      headers: auth.headers,
      reasoning: model.reasoning ? "high" : undefined,
      onPayload: buildPayloadPatcher(model),
    },
  );

  if (response.stopReason === "error") {
    throw new Error(`Summarization failed: ${response.errorMessage || "Unknown error"}`);
  }

  return textFromAssistantMessage(response);
}

export async function generateStructuredSummaryInChunks(
  chunks: any[][],
  model: any,
  reserveTokens: number,
  auth: { apiKey?: string; headers?: Record<string, string> },
  signal?: AbortSignal,
  customInstructions?: string,
  previousSummary?: string,
): Promise<{ summary: string; chunkCount: number }> {
  let runningSummary = previousSummary;

  for (const chunk of chunks) {
    runningSummary = await generateStructuredSummary(
      chunk,
      model,
      reserveTokens,
      auth,
      signal,
      customInstructions,
      runningSummary,
    );
  }

  return {
    summary: runningSummary?.trim() ?? "",
    chunkCount: chunks.length,
  };
}