import { existsSync, mkdirSync, readFileSync, writeFileSync } from "node:fs";
import { dirname, resolve } from "node:path";
import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";

const PREFERENCES_RELATIVE_PATH = "preference.md";
const LONG_TERM_SIGNAL_RE = /(以后|默认|统一|后续|今后|记住|长期|一直|都按这个)/;
const TEMPORARY_SIGNAL_RE = /(这次|本次|先|暂时|今天|当前)/;
const SENTENCE_SPLIT_RE = /[\n。！？；;]+/;

const SECTION_ORDER = ["回复偏好", "项目约定", "安全约束", "其他偏好"] as const;

export type PreferenceSection = (typeof SECTION_ORDER)[number];

export type PreferenceRecord = {
  key: string;
  section: PreferenceSection;
  label: string;
  value: string;
};

export default function preferenceMemoryExtension(pi: ExtensionAPI) {
  pi.on("before_agent_start", async (event, ctx) => {
    const preferencesPath = getPreferencesFilePath(ctx.cwd);
    const content = readInjectablePreferences(preferencesPath);

    return {
      systemPrompt: `${event.systemPrompt}

## Preference Persistence

This runtime can persist a small set of stable user or project preferences to workspace/preference.md.
Only persist preferences when the user clearly expresses a long-term default or rule, for example using phrases like "以后", "默认", "统一", "后续", or "记住".
Do not claim that you are unable to persist preferences in this runtime.
Do not say a preference was saved unless it actually matched the rules and was written.

${content ? `The following preferences were already persisted and should be treated as defaults unless the user overrides them for the current request:

${content}` : "No preferences have been persisted yet."}`
    };
  });

  pi.on("agent_end", async (event, ctx) => {
    const userPrompt = extractUserPrompt(event.messages);
    if (!userPrompt) {
      return;
    }

    const extracted = extractPreferencesFromText(userPrompt);
    if (extracted.length === 0) {
      return;
    }

    const preferencesPath = getPreferencesFilePath(ctx.cwd);
    const current = existsSync(preferencesPath) ? readFileSync(preferencesPath, "utf-8") : "";
    const next = upsertPreferencesDocument(current, extracted);
    if (normalizeDocument(current) === normalizeDocument(next)) {
      return;
    }

    mkdirSync(dirname(preferencesPath), { recursive: true });
    writeFileSync(preferencesPath, next, "utf-8");
  });
}

export function getPreferencesFilePath(cwd: string): string {
  return resolve(cwd, PREFERENCES_RELATIVE_PATH);
}

export function readInjectablePreferences(filePath: string): string | undefined {
  if (!existsSync(filePath)) {
    return undefined;
  }

  const content = readFileSync(filePath, "utf-8").trim();
  if (!content.includes("<!-- pref:")) {
    return undefined;
  }

  return content;
}

export function extractPreferencesFromText(text: string): PreferenceRecord[] {
  const preferences = new Map<string, PreferenceRecord>();

  for (const sentence of splitIntoSentences(text)) {
    const normalized = sentence.trim();
    if (!normalized) {
      continue;
    }

    const hasLongTermSignal = LONG_TERM_SIGNAL_RE.test(normalized);
    const hasTemporaryOnlySignal = TEMPORARY_SIGNAL_RE.test(normalized) && !hasLongTermSignal;
    if (hasTemporaryOnlySignal) {
      continue;
    }

    addPreference(preferences, detectReplyLanguage(normalized, hasLongTermSignal));
    addPreference(preferences, detectReplyStyle(normalized, hasLongTermSignal));
    addPreference(preferences, detectPackageManager(normalized, hasLongTermSignal));
    addPreference(preferences, detectRuntimeResourceDir(normalized, hasLongTermSignal));
    addPreference(preferences, detectDeleteConstraint(normalized, hasLongTermSignal));
  }

  return Array.from(preferences.values());
}

export function upsertPreferencesDocument(document: string, preferences: PreferenceRecord[]): string {
  const records = parseManagedPreferences(document);
  for (const preference of preferences) {
    records.set(preference.key, preference);
  }
  return renderPreferencesDocument(Array.from(records.values()));
}

export function parseManagedPreferences(document: string): Map<string, PreferenceRecord> {
  const records = new Map<string, PreferenceRecord>();
  const lines = document.split(/\r?\n/);

  for (const line of lines) {
    const match = line.match(/<!--\s*pref:([a-z0-9-]+)\s*-->\s*([^：:]+)[：:]\s*(.+)$/i);
    if (!match) {
      continue;
    }

    const key = match[1].trim();
    const label = match[2].trim();
    const value = match[3].trim();
    const section = inferSectionFromKey(key);
    if (!section || !label || !value) {
      continue;
    }

    records.set(key, { key, section, label, value });
  }

  return records;
}

export function renderPreferencesDocument(preferences: PreferenceRecord[]): string {
  const lines: string[] = [
    "# Preferences",
    "",
    "<!-- Managed by preference-memory extension. -->",
    ""
  ];

  const grouped = new Map<PreferenceSection, PreferenceRecord[]>();
  for (const section of SECTION_ORDER) {
    grouped.set(section, []);
  }

  for (const preference of preferences) {
    grouped.get(preference.section)?.push(preference);
  }

  for (const section of SECTION_ORDER) {
    lines.push(`## ${section}`, "");
    const items = (grouped.get(section) ?? []).sort((left, right) => left.label.localeCompare(right.label));
    for (const item of items) {
      lines.push(`- <!-- pref:${item.key} --> ${item.label}：${item.value}`);
    }
    lines.push("");
  }

  return `${trimTrailingBlankLines(lines).join("\n")}\n`;
}

export function extractUserPrompt(messages: unknown[]): string | undefined {
  for (let index = messages.length - 1; index >= 0; index -= 1) {
    const text = extractMessageText(messages[index], "user");
    if (text) {
      return text;
    }
  }

  return undefined;
}

function extractMessageText(message: unknown, role: string): string | undefined {
  if (!message || typeof message !== "object") {
    return undefined;
  }

  if ((message as { role?: unknown }).role !== role) {
    return undefined;
  }

  const content = (message as { content?: unknown }).content;
  if (typeof content === "string") {
    return content.trim() || undefined;
  }

  if (!Array.isArray(content)) {
    return undefined;
  }

  const text = content
    .flatMap((part) => {
      if (!part || typeof part !== "object") {
        return [];
      }

      const typedPart = part as { type?: unknown; text?: unknown };
      return typedPart.type === "text" && typeof typedPart.text === "string" ? [typedPart.text] : [];
    })
    .join("")
    .trim();

  return text || undefined;
}

function splitIntoSentences(text: string): string[] {
  return text
    .split(SENTENCE_SPLIT_RE)
    .map((sentence) => sentence.trim())
    .filter(Boolean);
}

function addPreference(
  records: Map<string, PreferenceRecord>,
  preference: PreferenceRecord | undefined
): void {
  if (!preference) {
    return;
  }

  records.set(preference.key, preference);
}

function detectReplyLanguage(sentence: string, hasLongTermSignal: boolean): PreferenceRecord | undefined {
  if (!hasLongTermSignal) {
    return undefined;
  }

  if (!/(回复|回答|输出|沟通|说明)/.test(sentence)) {
    return undefined;
  }

  if (/中文/.test(sentence)) {
    return createPreference("reply-language", "回复偏好", "回复语言", "中文");
  }

  if (/(英文|英语|english)/i.test(sentence)) {
    return createPreference("reply-language", "回复偏好", "回复语言", "英文");
  }

  return undefined;
}

function detectReplyStyle(sentence: string, hasLongTermSignal: boolean): PreferenceRecord | undefined {
  if (!hasLongTermSignal) {
    return undefined;
  }

  if (!/(回复|回答|输出|风格|说明|解释)/.test(sentence)) {
    return undefined;
  }

  if (/(简洁|简短|精简)/.test(sentence)) {
    return createPreference("reply-style", "回复偏好", "回复风格", "简洁");
  }

  if (/(详细|展开|细一点)/.test(sentence)) {
    return createPreference("reply-style", "回复偏好", "回复风格", "详细");
  }

  return undefined;
}

function detectPackageManager(sentence: string, hasLongTermSignal: boolean): PreferenceRecord | undefined {
  if (!hasLongTermSignal) {
    return undefined;
  }

  const packageManager = sentence.match(/\b(pnpm|npm|yarn|bun)\b/i)?.[1]?.toLowerCase();
  if (!packageManager) {
    return undefined;
  }

  if (!/(统一|默认|以后|后续|项目|依赖|安装)/.test(sentence)) {
    return undefined;
  }

  return createPreference("package-manager", "项目约定", "包管理器", packageManager);
}

function detectRuntimeResourceDir(sentence: string, hasLongTermSignal: boolean): PreferenceRecord | undefined {
  if (!hasLongTermSignal) {
    return undefined;
  }

  if (!/(写到|放到|创建在|存到|记录到)/.test(sentence)) {
    return undefined;
  }

  if (!/(workspace\/\.pi|\.pi\b)/.test(sentence)) {
    return undefined;
  }

  return createPreference("runtime-resource-dir", "项目约定", "运行态资源目录", "workspace/.pi");
}

function detectDeleteConstraint(sentence: string, hasLongTermSignal: boolean): PreferenceRecord | undefined {
  const hasStrongConstraint = /(没有我确认|未经确认|没确认|记住)/.test(sentence);
  if (!hasLongTermSignal && !hasStrongConstraint) {
    return undefined;
  }

  if (!/(删|删除|rm)/.test(sentence)) {
    return undefined;
  }

  if (!/(文件|目录|代码|资源)/.test(sentence)) {
    return undefined;
  }

  return createPreference("delete-confirmation", "安全约束", "删除约束", "未经确认不要删除文件");
}

function createPreference(
  key: string,
  section: PreferenceSection,
  label: string,
  value: string
): PreferenceRecord {
  return { key, section, label, value };
}

function inferSectionFromKey(key: string): PreferenceSection | undefined {
  if (key.startsWith("reply-")) {
    return "回复偏好";
  }

  if (key === "package-manager" || key === "runtime-resource-dir") {
    return "项目约定";
  }

  if (key === "delete-confirmation") {
    return "安全约束";
  }

  return "其他偏好";
}

function trimTrailingBlankLines(lines: string[]): string[] {
  const trimmed = [...lines];
  while (trimmed.length > 0 && trimmed[trimmed.length - 1].trim() === "") {
    trimmed.pop();
  }
  return trimmed;
}

function normalizeDocument(document: string): string {
  return document.replace(/\r\n/g, "\n").trim();
}