import { describe, it, expect, vi } from 'vitest';
import { AgentService, type VirtualAgentInput } from '../src/services/agent.service.js';
import { VirtualLlmService } from '../src/services/virtual-llm.service.js';
import type { IAgentRepository } from '../src/repositories/agent.repository.js';
import type { ILlmRepository } from '../src/repositories/llm.repository.js';
import type { LlmService } from '../src/services/llm.service.js';
import type { ProjectService } from '../src/services/project.service.js';
import type { Agent, Llm } from '@prisma/client';

/**
 * v3 Stage 2 — virtual-agent lifecycle methods on AgentService and the
 * cascade callbacks wired into VirtualLlmService.gcSweep / heartbeat /
 * unbindSession. Mirrors the shape of virtual-llm-service.test.ts but
 * focused on the agent-side state machine + the Llm→Agent cascade.
 */

const NOW = new Date();

function makeAgent(overrides: Partial<Agent> = {}): Agent {
  return {
    id: `agent-${Math.random().toString(36).slice(2, 8)}`,
    name: 'fake-agent',
    description: '',
    systemPrompt: '',
    llmId: 'llm-1',
    projectId: null,
    defaultPersonalityId: null,
    proxyModelName: null,
    defaultParams: {} as Agent['defaultParams'],
    extras: {} as Agent['extras'],
    kind: 'virtual',
    providerSessionId: 'sess-1',
    lastHeartbeatAt: NOW,
    status: 'active',
    inactiveSince: null,
    ownerId: 'owner-1',
    version: 1,
    createdAt: NOW,
    updatedAt: NOW,
    ...overrides,
  };
}

function makeLlm(overrides: Partial<Llm> = {}): Llm {
  return {
    id: `llm-${Math.random().toString(36).slice(2, 8)}`,
    name: 'vllm-local',
    type: 'openai',
    model: 'm',
    url: '',
    tier: 'fast',
    description: '',
    apiKeySecretId: null,
    apiKeySecretKey: null,
    extraConfig: {} as Llm['extraConfig'],
    kind: 'virtual',
    providerSessionId: 'sess-1',
    lastHeartbeatAt: NOW,
    status: 'active',
    inactiveSince: null,
    version: 1,
    createdAt: NOW,
    updatedAt: NOW,
    ...overrides,
  };
}

function mockAgentRepo(initial: Agent[] = []): IAgentRepository {
  const rows = new Map<string, Agent>(initial.map((r) => [r.id, r]));
  let counter = rows.size;
  return {
    findAll: vi.fn(async () => [...rows.values()]),
    findById: vi.fn(async (id: string) => rows.get(id) ?? null),
    findByName: vi.fn(async (name: string) => {
      for (const r of rows.values()) if (r.name === name) return r;
      return null;
    }),
    findByProjectId: vi.fn(async () => []),
    findBySessionId: vi.fn(async (sid: string) =>
      [...rows.values()].filter((r) => r.providerSessionId === sid)),
    findByLlmId: vi.fn(async (llmId: string) =>
      [...rows.values()].filter((r) => r.llmId === llmId)),
    findStaleVirtuals: vi.fn(async (cutoff: Date) =>
      [...rows.values()].filter((r) =>
        r.kind === 'virtual'
        && r.status === 'active'
        && r.lastHeartbeatAt !== null
        && r.lastHeartbeatAt < cutoff)),
    findExpiredInactives: vi.fn(async (cutoff: Date) =>
      [...rows.values()].filter((r) =>
        r.kind === 'virtual'
        && r.status === 'inactive'
        && r.inactiveSince !== null
        && r.inactiveSince < cutoff)),
    create: vi.fn(async (data) => {
      counter += 1;
      const row = makeAgent({
        id: `agent-${String(counter)}`,
        name: data.name,
        description: data.description ?? '',
        systemPrompt: data.systemPrompt ?? '',
        llmId: data.llmId,
        projectId: data.projectId ?? null,
        kind: data.kind ?? 'public',
        providerSessionId: data.providerSessionId ?? null,
        status: data.status ?? 'active',
        lastHeartbeatAt: data.lastHeartbeatAt ?? null,
        inactiveSince: data.inactiveSince ?? null,
        ownerId: data.ownerId,
      });
      rows.set(row.id, row);
      return row;
    }),
    update: vi.fn(async (id, data) => {
      const existing = rows.get(id);
      if (!existing) throw new Error('not found');
      const next: Agent = {
        ...existing,
        ...(data.description !== undefined ? { description: data.description } : {}),
        ...(data.systemPrompt !== undefined ? { systemPrompt: data.systemPrompt } : {}),
        ...(data.llmId !== undefined ? { llmId: data.llmId } : {}),
        ...(data.projectId !== undefined ? { projectId: data.projectId } : {}),
        ...(data.kind !== undefined ? { kind: data.kind } : {}),
        ...(data.providerSessionId !== undefined ? { providerSessionId: data.providerSessionId } : {}),
        ...(data.status !== undefined ? { status: data.status } : {}),
        ...(data.lastHeartbeatAt !== undefined ? { lastHeartbeatAt: data.lastHeartbeatAt } : {}),
        ...(data.inactiveSince !== undefined ? { inactiveSince: data.inactiveSince } : {}),
      };
      rows.set(id, next);
      return next;
    }),
    delete: vi.fn(async (id: string) => { rows.delete(id); }),
  };
}

function mockLlms(): LlmService {
  return {
    getById: vi.fn(async (id: string) => ({ id, name: 'vllm-local', type: 'openai', model: 'm', kind: 'virtual', status: 'active' })),
    getByName: vi.fn(async (name: string) => ({ id: 'llm-1', name, type: 'openai', model: 'm', kind: 'virtual', status: 'active' })),
  } as unknown as LlmService;
}

function mockProjects(opts: { existingNames?: string[] } = {}): ProjectService & { _created: string[] } {
  const existing = new Set(opts.existingNames ?? ['mcpctl-dev']);
  const created: string[] = [];
  // Reverse-lookup id → name so toView's getById finds the project we
  // ensured/created (instead of always returning a hardcoded name).
  const idToName = new Map<string, string>();
  const remember = (name: string): { id: string; name: string } => {
    const id = `proj-${name}`;
    idToName.set(id, name);
    return { id, name };
  };
  // Pre-seed the lookup so existing projects round-trip correctly.
  for (const n of existing) remember(n);
  const svc = {
    _created: created,
    getById: vi.fn(async (id: string) => ({ id, name: idToName.get(id) ?? 'mcpctl-dev' })),
    resolveAndGet: vi.fn(async (idOrName: string) => {
      if (!existing.has(idOrName)) {
        const err = new Error(`Project not found: ${idOrName}`);
        err.name = 'NotFoundError';
        throw err;
      }
      return remember(idOrName);
    }),
    // v6: ensureByName auto-creates when missing.
    ensureByName: vi.fn(async (name: string, ownerId: string, _opts?: { auditNote?: string }) => {
      if (existing.has(name)) return remember(name);
      created.push(name);
      existing.add(name);
      const row = remember(name);
      return { ...row, ownerId };
    }),
  };
  return svc as unknown as ProjectService & { _created: string[] };
}

describe('AgentService — virtual-agent lifecycle (v3 Stage 2)', () => {
  it('registerVirtualAgents inserts new rows with kind=virtual / status=active', async () => {
    const repo = mockAgentRepo();
    const svc = new AgentService(repo, mockLlms(), mockProjects());
    const inputs: VirtualAgentInput[] = [
      { name: 'local-coder', llmName: 'vllm-local', description: 'd', systemPrompt: 's' },
    ];
    const out = await svc.registerVirtualAgents('sess-1', inputs, 'owner-1');
    expect(out).toHaveLength(1);
    expect(out[0]!.kind).toBe('virtual');
    expect(out[0]!.status).toBe('active');
  });

  it('registerVirtualAgents reuses an existing row from the same session (sticky reconnect)', async () => {
    const existing = makeAgent({ name: 'local-coder', providerSessionId: 'sess-1', status: 'inactive', inactiveSince: NOW });
    const repo = mockAgentRepo([existing]);
    const svc = new AgentService(repo, mockLlms(), mockProjects());
    const out = await svc.registerVirtualAgents(
      'sess-1',
      [{ name: 'local-coder', llmName: 'vllm-local' }],
      'owner-1',
    );
    expect(out[0]!.id).toBe(existing.id);
    expect(out[0]!.status).toBe('active');
  });

  it('registerVirtualAgents adopts an inactive virtual from a different session', async () => {
    const existing = makeAgent({
      name: 'local-coder', providerSessionId: 'old-session',
      status: 'inactive', inactiveSince: NOW,
    });
    const repo = mockAgentRepo([existing]);
    const svc = new AgentService(repo, mockLlms(), mockProjects());
    const out = await svc.registerVirtualAgents(
      'new-session',
      [{ name: 'local-coder', llmName: 'vllm-local' }],
      'owner-1',
    );
    expect(out[0]!.id).toBe(existing.id);
    expect(out[0]!.status).toBe('active');
  });

  it('registerVirtualAgents auto-creates a missing project (v6 Stage 2)', async () => {
    // Pre-v6 the publisher had to ask an admin to create the project
    // first; otherwise registerVirtualAgents 404'd. v6: ensureByName
    // creates with sensible defaults and the publishing user as owner.
    const repo = mockAgentRepo();
    const projects = mockProjects({ existingNames: [] });
    const svc = new AgentService(repo, mockLlms(), projects);
    const out = await svc.registerVirtualAgents(
      'sess-1',
      [{ name: 'team-coder', llmName: 'vllm-local', project: 'platform' }],
      'owner-platform',
    );
    expect(out).toHaveLength(1);
    expect(out[0]!.project?.name).toBe('platform');
    expect(projects._created).toEqual(['platform']);
  });

  it('registerVirtualAgents reuses an existing project without re-creating it (v6 Stage 2)', async () => {
    // ensureByName must be idempotent — second publisher landing on
    // the same project name doesn't try to re-create (would 409).
    const repo = mockAgentRepo();
    const projects = mockProjects({ existingNames: ['platform'] });
    const svc = new AgentService(repo, mockLlms(), projects);
    await svc.registerVirtualAgents(
      'sess-2',
      [{ name: 'team-reviewer', llmName: 'vllm-local', project: 'platform' }],
      'owner-other',
    );
    expect(projects._created).toEqual([]);
  });

  it('registerVirtualAgents refuses to overwrite a public agent (409)', async () => {
    const repo = mockAgentRepo([makeAgent({ name: 'reviewer', kind: 'public', providerSessionId: null })]);
    const svc = new AgentService(repo, mockLlms(), mockProjects());
    await expect(svc.registerVirtualAgents(
      'sess-x',
      [{ name: 'reviewer', llmName: 'vllm-local' }],
      'owner-1',
    )).rejects.toThrow(/Cannot publish over public Agent/);
  });

  it('registerVirtualAgents refuses if another active session owns the name', async () => {
    const repo = mockAgentRepo([makeAgent({ name: 'local-coder', providerSessionId: 'other', status: 'active' })]);
    const svc = new AgentService(repo, mockLlms(), mockProjects());
    await expect(svc.registerVirtualAgents(
      'mine',
      [{ name: 'local-coder', llmName: 'vllm-local' }],
      'owner-1',
    )).rejects.toThrow(/already active under a different session/);
  });

  it('heartbeatVirtualAgents bumps + revives inactive', async () => {
    const past = new Date(Date.now() - 5_000);
    const a = makeAgent({ name: 'a', providerSessionId: 'sess', status: 'inactive', lastHeartbeatAt: past, inactiveSince: past });
    const repo = mockAgentRepo([a]);
    const svc = new AgentService(repo, mockLlms(), mockProjects());
    await svc.heartbeatVirtualAgents('sess');
    const row = await repo.findByName('a');
    expect(row?.status).toBe('active');
    expect(row?.inactiveSince).toBeNull();
    expect(row!.lastHeartbeatAt!.getTime()).toBeGreaterThan(past.getTime());
  });

  it('markVirtualAgentsInactiveBySession flips owned actives to inactive', async () => {
    const repo = mockAgentRepo([
      makeAgent({ name: 'a', providerSessionId: 'sess' }),
      makeAgent({ name: 'b', providerSessionId: 'sess' }),
      makeAgent({ name: 'c', providerSessionId: 'other' }),
    ]);
    const svc = new AgentService(repo, mockLlms(), mockProjects());
    await svc.markVirtualAgentsInactiveBySession('sess');
    expect((await repo.findByName('a'))?.status).toBe('inactive');
    expect((await repo.findByName('b'))?.status).toBe('inactive');
    expect((await repo.findByName('c'))?.status).toBe('active');
  });

  it('deleteVirtualAgentsForLlm deletes only virtuals pinned to that Llm', async () => {
    const repo = mockAgentRepo([
      makeAgent({ name: 'v-1', llmId: 'doomed', kind: 'virtual' }),
      makeAgent({ name: 'v-2', llmId: 'doomed', kind: 'virtual' }),
      makeAgent({ name: 'pub-1', llmId: 'doomed', kind: 'public', providerSessionId: null }),
      makeAgent({ name: 'v-other', llmId: 'safe', kind: 'virtual' }),
    ]);
    const svc = new AgentService(repo, mockLlms(), mockProjects());
    const deleted = await svc.deleteVirtualAgentsForLlm('doomed');
    expect(deleted).toBe(2);
    expect(await repo.findByName('v-1')).toBeNull();
    expect(await repo.findByName('v-2')).toBeNull();
    expect(await repo.findByName('pub-1')).not.toBeNull();
    expect(await repo.findByName('v-other')).not.toBeNull();
  });

  it('gcSweepVirtualAgents flips heartbeat-stale + deletes 4h-old inactive', async () => {
    const long = new Date(Date.now() - 5 * 60 * 1000);    // 5 min ago, past 90s cutoff
    const ancient = new Date(Date.now() - 5 * 60 * 60 * 1000); // 5 h ago, past 4h cutoff
    const repo = mockAgentRepo([
      makeAgent({ name: 'stale', providerSessionId: 'a', status: 'active', lastHeartbeatAt: long }),
      makeAgent({ name: 'old', providerSessionId: 'b', status: 'inactive', inactiveSince: ancient }),
      makeAgent({ name: 'pub', providerSessionId: null, kind: 'public' }),
    ]);
    const svc = new AgentService(repo, mockLlms(), mockProjects());
    const r = await svc.gcSweepVirtualAgents();
    expect(r.markedInactive).toBe(1);
    expect(r.deleted).toBe(1);
    expect((await repo.findByName('stale'))?.status).toBe('inactive');
    expect(await repo.findByName('old')).toBeNull();
    expect(await repo.findByName('pub')).not.toBeNull();
  });
});

describe('VirtualLlmService cascade through AgentService (v3 Stage 2)', () => {
  function mockLlmRepo(initial: Llm[] = []): ILlmRepository {
    const rows = new Map<string, Llm>(initial.map((r) => [r.id, r]));
    let counter = rows.size;
    return {
      findAll: vi.fn(async () => [...rows.values()]),
      findById: vi.fn(async (id: string) => rows.get(id) ?? null),
      findByName: vi.fn(async (name: string) => {
        for (const r of rows.values()) if (r.name === name) return r;
        return null;
      }),
      findByTier: vi.fn(async () => []),
      findBySessionId: vi.fn(async (sid: string) =>
        [...rows.values()].filter((r) => r.providerSessionId === sid)),
      findStaleVirtuals: vi.fn(async (cutoff: Date) =>
        [...rows.values()].filter((r) =>
          r.kind === 'virtual'
          && r.status === 'active'
          && r.lastHeartbeatAt !== null
          && r.lastHeartbeatAt < cutoff)),
      findExpiredInactives: vi.fn(async (cutoff: Date) =>
        [...rows.values()].filter((r) =>
          r.kind === 'virtual'
          && r.status === 'inactive'
          && r.inactiveSince !== null
          && r.inactiveSince < cutoff)),
      create: vi.fn(async (data) => {
        counter += 1;
        const row = makeLlm({ id: `llm-${String(counter)}`, name: data.name, type: data.type });
        rows.set(row.id, row);
        return row;
      }),
      update: vi.fn(async (id, data) => {
        const existing = rows.get(id);
        if (!existing) throw new Error('not found');
        const next: Llm = { ...existing, ...data } as Llm;
        rows.set(id, next);
        return next;
      }),
      delete: vi.fn(async (id: string) => { rows.delete(id); }),
    };
  }

  it('unbindSession cascades to mark virtual agents inactive', async () => {
    const llmRepo = mockLlmRepo([makeLlm({ name: 'vllm-local', providerSessionId: 'sess' })]);
    const agentRepo = mockAgentRepo([
      makeAgent({ name: 'local-coder', providerSessionId: 'sess' }),
    ]);
    const agents = new AgentService(agentRepo, mockLlms(), mockProjects());
    const svc = new VirtualLlmService(llmRepo, agents);
    await svc.unbindSession('sess');
    expect((await agentRepo.findByName('local-coder'))?.status).toBe('inactive');
  });

  it('gcSweep deletes virtual agents BEFORE their pinned virtual Llm', async () => {
    const ancient = new Date(Date.now() - 5 * 60 * 60 * 1000);
    const llmRepo = mockLlmRepo([makeLlm({
      id: 'doomed-llm', name: 'vllm-local', providerSessionId: 'sess',
      status: 'inactive', inactiveSince: ancient,
    })]);
    const agentRepo = mockAgentRepo([
      makeAgent({ name: 'pinned', providerSessionId: 'sess', llmId: 'doomed-llm', status: 'inactive', inactiveSince: ancient }),
    ]);
    const agents = new AgentService(agentRepo, mockLlms(), mockProjects());
    const svc = new VirtualLlmService(llmRepo, agents);
    const r = await svc.gcSweep();
    expect(r.deleted).toBeGreaterThanOrEqual(2); // 1 agent + 1 llm
    expect(await llmRepo.findByName('vllm-local')).toBeNull();
    expect(await agentRepo.findByName('pinned')).toBeNull();
  });

  it('gcSweep defensive cascade: still drops the agent when its heartbeat lagged the Llm', async () => {
    // The Llm is past the 4h cutoff. The agent is inactive but only
    // 1h old — wouldn't be GC'd by gcSweepVirtualAgents on its own.
    // The defensive cascade in gcSweep deletes it anyway because the
    // Restrict FK would otherwise block the Llm delete.
    const ancient = new Date(Date.now() - 5 * 60 * 60 * 1000);
    const recent = new Date(Date.now() - 1 * 60 * 60 * 1000);
    const llmRepo = mockLlmRepo([makeLlm({
      id: 'doomed-llm', name: 'vllm-local', providerSessionId: 'sess',
      status: 'inactive', inactiveSince: ancient,
    })]);
    const agentRepo = mockAgentRepo([
      makeAgent({ name: 'pinned', providerSessionId: 'sess', llmId: 'doomed-llm', status: 'inactive', inactiveSince: recent }),
    ]);
    const agents = new AgentService(agentRepo, mockLlms(), mockProjects());
    const svc = new VirtualLlmService(llmRepo, agents);
    await svc.gcSweep();
    expect(await llmRepo.findByName('vllm-local')).toBeNull();
    expect(await agentRepo.findByName('pinned')).toBeNull();
  });

  it('heartbeat cascades to bump owned virtual agents', async () => {
    const past = new Date(Date.now() - 10_000);
    const llmRepo = mockLlmRepo([makeLlm({ name: 'vllm-local', providerSessionId: 'sess', lastHeartbeatAt: past })]);
    const agentRepo = mockAgentRepo([makeAgent({ name: 'local-coder', providerSessionId: 'sess', lastHeartbeatAt: past })]);
    const agents = new AgentService(agentRepo, mockLlms(), mockProjects());
    const svc = new VirtualLlmService(llmRepo, agents);
    await svc.heartbeat('sess');
    const a = await agentRepo.findByName('local-coder');
    expect(a!.lastHeartbeatAt!.getTime()).toBeGreaterThan(past.getTime());
  });
});
