import request from "supertest";
import { afterEach, describe, expect, it, vi } from "vitest";
import {
  INVESTIGATION_MAX_ACTIVE_PROBLEM_UNITS,
  INVESTIGATION_MAX_FOCUS_PROBLEM_TYPES,
  INVESTIGATION_MAX_PROBLEM_UNIT_BACKLINKS,
  INVESTIGATION_MAX_RESOLVED_PROBLEM_UNITS
} from "../src/types/stage2ProblemUnits";

const FLAG_KEYS = [
  "FEATURE_ASSISTANT_INVESTIGATION_STATE_V1",
  "FEATURE_ASSISTANT_STATE_FOLLOWUP_BINDING_V1",
  "FEATURE_ASSISTANT_PROBLEM_UNITS_V1",
  "FEATURE_ASSISTANT_PROBLEM_UNIT_CONTINUITY_V1"
] as const;

const ORIGINAL_FLAGS: Record<string, string | undefined> = Object.fromEntries(
  FLAG_KEYS.map((key) => [key, process.env[key]])
);

function restoreFlags(): void {
  for (const key of FLAG_KEYS) {
    const original = ORIGINAL_FLAGS[key];
    if (original === undefined) {
      delete process.env[key];
    } else {
      process.env[key] = original;
    }
  }
}

async function createAppWithWave4Flags() {
  process.env.FEATURE_ASSISTANT_INVESTIGATION_STATE_V1 = "1";
  process.env.FEATURE_ASSISTANT_STATE_FOLLOWUP_BINDING_V1 = "1";
  process.env.FEATURE_ASSISTANT_PROBLEM_UNITS_V1 = "1";
  process.env.FEATURE_ASSISTANT_PROBLEM_UNIT_CONTINUITY_V1 = "1";
  vi.resetModules();
  const { createApp } = await import("../src/server");
  return createApp();
}

describe.sequential("assistant problem-unit continuity state", () => {
  afterEach(() => {
    restoreFlags();
    vi.resetModules();
  });

  it("stores bounded problem_unit_state in investigation snapshot and session state", async () => {
    const app = await createAppWithWave4Flags();
    const sessionId = `asst-wave4-state-${Date.now()}`;

    const first = await request(app).post("/api/assistant/message").send({
      session_id: sessionId,
      useMock: true,
      promptVersion: "normalizer_v2_0_2",
      user_message: "Разложи цепочку документов и оплат по контрагентам за 2020-06, где есть разрыв закрытия."
    });

    expect(first.status).toBe(200);
    expect(first.body.debug?.investigation_state_snapshot?.problem_unit_state).toBeTruthy();

    const second = await request(app).post("/api/assistant/message").send({
      session_id: sessionId,
      useMock: true,
      promptVersion: "normalizer_v2_0_2",
      user_message: "И по тому же кейсу уточни, что влияет на закрытие периода."
    });

    expect(second.status).toBe(200);
    const problemState = second.body.debug?.investigation_state_snapshot?.problem_unit_state;
    expect(problemState).toBeTruthy();
    expect(problemState.active_problem_units.length).toBeLessThanOrEqual(INVESTIGATION_MAX_ACTIVE_PROBLEM_UNITS);
    expect(problemState.resolved_problem_units.length).toBeLessThanOrEqual(INVESTIGATION_MAX_RESOLVED_PROBLEM_UNITS);
    expect(problemState.problem_unit_backlinks.length).toBeLessThanOrEqual(INVESTIGATION_MAX_PROBLEM_UNIT_BACKLINKS);
    expect(problemState.focus_problem_types.length).toBeLessThanOrEqual(INVESTIGATION_MAX_FOCUS_PROBLEM_TYPES);

    const sessionResponse = await request(app).get(`/api/assistant/session/${sessionId}`);
    expect(sessionResponse.status).toBe(200);
    const sessionProblemState = sessionResponse.body.session?.investigation_state?.problem_unit_state;
    expect(sessionProblemState).toBeTruthy();
    expect(sessionProblemState.active_problem_units.length).toBeLessThanOrEqual(INVESTIGATION_MAX_ACTIVE_PROBLEM_UNITS);
    expect(sessionProblemState.resolved_problem_units.length).toBeLessThanOrEqual(INVESTIGATION_MAX_RESOLVED_PROBLEM_UNITS);
    expect(sessionProblemState.problem_unit_backlinks.length).toBeLessThanOrEqual(INVESTIGATION_MAX_PROBLEM_UNIT_BACKLINKS);
    expect(sessionProblemState.focus_problem_types.length).toBeLessThanOrEqual(INVESTIGATION_MAX_FOCUS_PROBLEM_TYPES);
  });
});