// ==UserScript==
// @name         NovelAI V5 多様性コントロール
// @namespace    local.novelai.v5-variety
// @version      1.0.2
// @description  V5の生成リクエストに多様性の値を追加。ON/OFF・数値を画面から設定できます。
// @match        https://novelai.net/*
// @run-at       document-start
// @sandbox      raw
// @grant        none
// @noframes
// ==/UserScript==

(() => {
  'use strict';

  const INSTALL_KEY = '__novelaiV5VarietyInstalled_v1';
  if (window[INSTALL_KEY]) return;
  window[INSTALL_KEY] = true;

  let panelHost = null;
  // Register at document-start, before the site's shortcut handlers. Outside a
  // shadow root, activeElement/target refer to the host DIV, not its input, so
  // the site's "ignore input fields" check does not protect this panel.
  for (const type of ['keydown', 'keypress', 'keyup']) {
    window.addEventListener(type, event => {
      if (panelHost && event.composedPath().includes(panelHost)) {
        event.stopImmediatePropagation();
        // Do NOT preventDefault: typing, selection, paste and Tab still work.
      }
    }, { capture: true });
  }

  const STORAGE_KEY = 'local.novelai.v5-variety.settings.v1';
  const MODELS = new Set([
    'nai-diffusion-5-curated',
    'nai-diffusion-5-curated-inpainting',
    'nai-diffusion-5-full',
    'nai-diffusion-5-full-inpainting',
    'nai-diffusion-5-full-medium',
    'nai-diffusion-5-full-medium-inpainting'
  ]);
  const defaults = { enabled: true, sigma: 58, collapsed: false, position: null };
  const validSigma = value => typeof value === 'number' && Number.isFinite(value) && value > 0;
  let settings = { ...defaults };
  try {
    const saved = JSON.parse(localStorage.getItem(STORAGE_KEY));
    if (saved && typeof saved === 'object') {
      if (typeof saved.enabled === 'boolean') settings.enabled = saved.enabled;
      if (validSigma(saved.sigma)) settings.sigma = saved.sigma;
      if (typeof saved.collapsed === 'boolean') settings.collapsed = saved.collapsed;
      if (saved.position && Number.isFinite(saved.position.x) && Number.isFinite(saved.position.y)) {
        settings.position = { x: saved.position.x, y: saved.position.y };
      }
    }
  } catch { /* Storage may be disabled; in-memory settings still work. */ }

  let status = 'まだ書き換えていません';
  let render = () => {};
  function report(message) {
    status = message;
    render();
  }
  function save() {
    try {
      localStorage.setItem(STORAGE_KEY, JSON.stringify(settings));
    } catch {
      report('設定を保存できません。このページ内では有効です。');
    }
    render();
  }

  // Patch only the official image-generation POST endpoints, including streams.
  function isTarget(input, init) {
    const request = input instanceof Request ? input : null;
    const method = String(init?.method ?? request?.method ?? 'GET').toUpperCase();
    if (method !== 'POST') return false;
    const url = new URL(request ? request.url : String(input), location.href);
    return url.origin === 'https://image.novelai.net'
      && /^\/ai\/generate-image(?:-stream)?\/?$/.test(url.pathname);
  }

  function patchJSON(text, sigma) {
    const payload = JSON.parse(text);
    if (!payload || !MODELS.has(payload.model)) return null;
    if (!payload.parameters || typeof payload.parameters !== 'object'
        || Array.isArray(payload.parameters)) {
      throw new Error('Missing parameters');
    }
    payload.parameters.skip_cfg_above_sigma = sigma;
    return JSON.stringify(payload);
  }

  async function patchBody(body, sigma) {
    if (body instanceof FormData) {
      const parts = body.getAll('request');
      if (parts.length !== 1) throw new Error('Unexpected request parts');
      const original = parts[0];
      const json = patchJSON(typeof original === 'string' ? original : await original.text(), sigma);
      if (json === null) return null;
      const copy = new FormData();
      for (const [name, value] of body.entries()) {
        if (name !== 'request') {
          // Keep filenames, MIME types, image bytes and repeated fields intact.
          copy.append(name, value);
        } else if (typeof original === 'string') {
          copy.append(name, json);
        } else {
          copy.append(name, new Blob([json], { type: original.type || 'application/json' }), original.name || 'blob');
        }
      }
      return { body: copy, multipart: true };
    }
    if (typeof body === 'string' || body instanceof Blob) {
      const json = patchJSON(typeof body === 'string' ? body : await body.text(), sigma);
      if (json === null) return null;
      return {
        body: typeof body === 'string' ? json : new Blob([json], { type: body.type || 'application/json' }),
        multipart: false
      };
    }
    throw new Error('Unsupported request body');
  }

  const originalFetch = window.fetch;
  async function varietyFetch(input, init) {
    let nextInit = init;
    // Snapshot settings before asynchronous Blob/Request reads.
    const { enabled, sigma } = settings;
    try {
      if (enabled && isTarget(input, init)) {
        const request = input instanceof Request ? input : null;
        const headers = new Headers(init?.headers ?? request?.headers);
        let body = init?.body;
        if (body == null && request) {
          const clone = request.clone();
          const type = request.headers.get('Content-Type') || '';
          body = /^multipart\/form-data\b/i.test(type) ? await clone.formData() : await clone.text();
        }
        const patched = await patchBody(body, sigma);
        if (patched) {
          // A new multipart body needs a new browser-generated boundary.
          if (patched.multipart) headers.delete('Content-Type');
          else if (!headers.has('Content-Type')) headers.set('Content-Type', 'application/json');
          nextInit = { ...init, headers, body: patched.body };
          report(`直近の書換え: ${sigma} (${new Date().toLocaleTimeString()})`);
        }
      }
    } catch {
      // Never log prompts/tokens, or retry a potentially billed generation.
      report('書換えできませんでした。元のリクエストを送信します。');
      nextInit = init;
    }
    // Keep network errors outside the rewriting catch: fetch is called ONCE.
    // Returning the original Response also leaves streaming untouched.
    return Reflect.apply(originalFetch, window, [input, nextInit]);
  }
  window.fetch = varietyFetch;

  function mount() {
    const host = document.createElement('div');
    panelHost = host;
    host.id = 'novelai-v5-variety-control';
    host.style.cssText = 'position:fixed;right:16px;bottom:16px;z-index:2147483647;';
    const root = host.attachShadow({ mode: 'open' });
    const style = document.createElement('style');
    style.textContent = `
      :host { color-scheme: dark; font: 13px/1.5 system-ui, sans-serif; }
      section { box-sizing: border-box; width: 272px; max-width: calc(100vw - 32px);
        max-height: calc(100dvh - 32px); overflow: auto;
        color: #f4f3ff; background: #242436; border: 1px solid #77768e;
        border-radius: 10px; padding: 12px; box-shadow: 0 4px 20px #0006; }
      header, label { display: flex; align-items: center; gap: 8px; }
      header { justify-content: space-between; }
      h2 { font-size: 14px; margin: 0; flex: 1; cursor: grab; user-select: none; touch-action: none; }
      h2[data-dragging] { cursor: grabbing; }
      button, input { font: inherit; } button { color: inherit; background: #3b3b53;
        border: 1px solid #77768e; border-radius: 5px; padding: 2px 9px; cursor: pointer; }
      label { margin-top: 10px; } input[type=checkbox] { accent-color: #b6a4ff; }
      input[type=number] { box-sizing: border-box; width: 104px; color: inherit;
        background: #171723; border: 1px solid #77768e; border-radius: 5px; padding: 4px 6px; }
      p { margin: 8px 0 0; font-size: 11px; color: #cbc9df; overflow-wrap: anywhere; }
      [hidden] { display: none !important; }
    `;
    const panel = document.createElement('section');
    panel.setAttribute('aria-label', 'V5 多様性コントロール');
    const header = document.createElement('header');
    const title = document.createElement('h2');
    title.title = 'ドラッグして移動（位置は保存されます）';
    const collapse = document.createElement('button');
    collapse.type = 'button';
    header.append(title, collapse);
    const controls = document.createElement('div');
    const toggleLabel = document.createElement('label');
    const toggle = document.createElement('input');
    toggle.type = 'checkbox';
    toggleLabel.append(toggle, '多様性を有効化');
    const numberLabel = document.createElement('label');
    const number = document.createElement('input');
    number.type = 'number';
    number.step = 'any';
    number.min = '0';
    number.required = true;
    number.value = String(settings.sigma);
    numberLabel.append('数値 (sigma)', number);
    const hint = document.createElement('p');
    hint.textContent = '0に近いほど暴れやすい。初期値: 58';
    const feedback = document.createElement('p');
    feedback.setAttribute('role', 'status');
    const lastRequest = document.createElement('p');
    lastRequest.setAttribute('aria-live', 'polite');
    controls.append(toggleLabel, numberLabel, hint, feedback, lastRequest);
    panel.append(header, controls);
    root.append(style, panel);
    document.body.append(host);

    function placePanel() {
      if (!settings.position || host.hidden) return;
      const margin = 8;
      const rect = host.getBoundingClientRect();
      const width = document.documentElement.clientWidth;
      const height = document.documentElement.clientHeight;
      const x = Math.min(Math.max(margin, settings.position.x), Math.max(margin, width - rect.width - margin));
      const y = Math.min(Math.max(margin, settings.position.y), Math.max(margin, height - rect.height - margin));
      host.style.right = 'auto';
      host.style.bottom = 'auto';
      host.style.left = `${x}px`;
      host.style.top = `${y}px`;
    }

    // Only the heading is a drag handle; inputs and the collapse button stay usable.
    let drag = null;
    title.addEventListener('pointerdown', event => {
      if (event.button !== 0 || !event.isPrimary) return;
      event.preventDefault();
      event.stopPropagation();
      const rect = host.getBoundingClientRect();
      drag = { id: event.pointerId, offsetX: event.clientX - rect.left, offsetY: event.clientY - rect.top };
      title.setPointerCapture(event.pointerId);
      title.setAttribute('data-dragging', '');
    });
    title.addEventListener('pointermove', event => {
      if (!drag || event.pointerId !== drag.id) return;
      event.stopPropagation();
      settings.position = { x: event.clientX - drag.offsetX, y: event.clientY - drag.offsetY };
      placePanel();
    });
    function finishDrag(event) {
      if (!drag || event.pointerId !== drag.id) return;
      event.stopPropagation();
      const id = drag.id;
      drag = null;
      title.removeAttribute('data-dragging');
      if (title.hasPointerCapture(id)) title.releasePointerCapture(id);
      const rect = host.getBoundingClientRect();
      settings.position = { x: rect.left, y: rect.top };
      save();
    }
    title.addEventListener('pointerup', finishDrag);
    title.addEventListener('pointercancel', finishDrag);
    title.addEventListener('lostpointercapture', finishDrag);
    window.addEventListener('resize', placePanel);
    // Keep the panel on-screen after expanding or when status text changes height.
    const panelResizeObserver = new ResizeObserver(placePanel);
    panelResizeObserver.observe(host);

    toggle.addEventListener('change', () => {
      settings.enabled = toggle.checked;
      save();
    });
    number.addEventListener('input', () => {
      const value = number.valueAsNumber;
      if (!validSigma(value)) {
        number.setCustomValidity('0より大きい有限の数値を入力してください。');
        feedback.textContent = `入力が無効です。適用値は ${settings.sigma} のままです。`;
        return;
      }
      number.setCustomValidity('');
      feedback.textContent = '';
      settings.sigma = value;
      save();
    });
    collapse.addEventListener('click', () => {
      settings.collapsed = !settings.collapsed;
      save();
    });
    render = () => {
      host.hidden = !/^\/image(?:\/|$)/.test(location.pathname);
      title.textContent = `V5 多様性: ${settings.enabled ? settings.sigma : 'OFF'}`;
      toggle.checked = settings.enabled;
      controls.hidden = settings.collapsed;
      collapse.textContent = settings.collapsed ? '開く' : '畳む';
      collapse.setAttribute('aria-expanded', String(!settings.collapsed));
      lastRequest.textContent = window.fetch === varietyFetch
        ? status : '送信処理が変更されました。ページを再読み込みしてください。';
      placePanel();
    };
    render();
    // NovelAI can navigate to /image without a document reload.
    window.setInterval(render, 1000);
  }
  if (document.readyState === 'loading') {
    document.addEventListener('DOMContentLoaded', mount, { once: true });
  } else {
    mount();
  }
})();