// agent routing extension for SootSimBridgeHost.
//
// owns the single FIFO reader per agent session (closes the
// "single-reader channel" gap that agent-sessions.ts explicitly warns
// about), and fans events out to every WS subscriber — CLI `sootsim
// agent watch`, electron main, and browser shells all consume the same
// stream.
//
// message shapes follow the existing ws-bridge convention:
//   client → daemon: { id, type: 'agent:…', …payload }
//   daemon → client: { id, result } | { id, error, code? }
// plus server-initiated pushes (no id):
//   { type: 'agent:event',         sessionId, event }
//   { type: 'agent:session-status', session }

import fs from 'node:fs'
import path from 'node:path'
import {
  scanDevServers,
  type DiscoveredServer,
} from '../../scripts/dev-server-scanner.ts'
import {
  AgentSessionError,
  endSession as endAgentSession,
  sendPrompt as sendAgentPrompt,
  startSession as startAgentSession,
  subscribeEvents as subscribeAgentEvents,
  transcriptPath,
  type Provider,
} from '../agent-sessions.ts'
import {
  deleteProject,
  findProjectById,
  findSessionById,
  getUserDataDir,
  listProjects,
  listSessions,
  recordTurnTelemetry,
  seedFromDemoAppRegistry,
  applySessionStatusPatch,
  updateSessionStatuses,
  upsertProject,
  type AgentSession,
  type AttachedProject,
} from '../attached-projects.ts'
import {
  sendPromptToTeamMachineSession,
  teamMachineSessionIdFromCliIdentity,
} from '../team-machine-prompt.ts'
import type { AgentEvent } from '../agent-events.ts'
import type { AgentPromptEnvelope } from '../agent-prompt.ts'
import type { WebSocket } from 'ws'

// ws.readyState constant — avoid importing the module just for this enum.
const WS_OPEN = 1
const SESSION_STATUS_PERSIST_DELAY_MS = 25

interface FanoutSubscription {
  unsubscribe: () => void
  refCount: number
}

interface PendingPromptEcho {
  sentAt: number
}

export interface AgentHostOptions {
  /** callback so hosts with a Contrast-env context (electron, dev middleware)
   *  can inject their own exclude list. defaults to the env-derived list
   *  the rest of the repo uses, so `rnx serve` standalone does the
   *  right thing. */
  getExcludePorts?: () => number[]
  /** returns the active CLI lease for a sim. used only for prompt routing to
   *  an external Team Machine agent that claimed the tab. */
  resolveCliLease?: (simId: string) => AgentCliLeaseSnapshot | null
}

export interface AgentCliLeaseSnapshot {
  kind: 'cli' | 'user-active'
  cliIdentityKey: string
  expiresAt: number
}

/** default excludes keep the dev-server scanner from hitting Contrast's own
 *  web / zero / r2 ports. mirrored from
 *  sootsim-engine/src/dev-scan-excludes.ts — duplicated rather than
 *  imported because sootsim is a dependency of sootsim-engine, not the
 *  other way around. */
function defaultExcludePorts(): number[] {
  return [
    Number(process.env.VITE_PORT_WEB || process.env.PORT || 3000),
    Number(process.env.VITE_PORT_ZERO || 7849),
    Number(process.env.VITE_PORT_R2 || 9500),
  ].filter((p) => Number.isFinite(p) && p > 0)
}

export class AgentHost {
  // live fan-out, keyed by sessionId. refcount tracks how many sockets have
  // subscribed; the FIFO reader stays open as long as refcount > 0.
  private subscriptions = new Map<string, FanoutSubscription>()
  // per-socket subscription set so ws.close can clean up without scanning
  // every session.
  private sessionsBySocket = new Map<WebSocket, Set<string>>()
  // every connected socket (regardless of role) — gets session-status
  // pushes so a CLI `sootsim agent sessions --watch` can react to state
  // changes driven by another client.
  private allSockets = new Set<WebSocket>()
  // agent wrappers also emit prompt-received from the FIFO reader, but the
  // shell/daemon already knows the user-facing display text when send-prompt
  // is accepted. emit the friendly prompt immediately and suppress the raw
  // wrapper echoes that follow a moment later.
  private pendingPromptEchoes = new Map<string, PendingPromptEcho[]>()
  // accepted prompts serialize inside the wrapper, so a second send while the
  // session is already working becomes queued follow-up work. keep a small
  // in-memory count so session status stays `working` until that backlog
  // truly drains.
  private pendingTurns = new Map<string, number>()
  // status fan-out is realtime state. keep the current value in memory and
  // batch its durable store write so an fsync cannot delay the websocket event
  // it describes. close() flushes the final batch synchronously.
  private pendingSessionStates = new Map<string, AgentSession>()
  private pendingSessionPatches = new Map<string, Partial<AgentSession>>()
  private sessionStatusPersistTimer: ReturnType<typeof setTimeout> | null = null
  private opts: AgentHostOptions

  constructor(opts: AgentHostOptions = {}) {
    this.opts = opts
  }

  registerSocket(ws: WebSocket): void {
    this.allSockets.add(ws)
  }

  unregisterSocket(ws: WebSocket): void {
    const sessions = this.sessionsBySocket.get(ws)
    if (sessions) {
      for (const sessionId of sessions) {
        this.decrementSubscription(sessionId)
      }
      this.sessionsBySocket.delete(ws)
    }
    this.allSockets.delete(ws)
  }

  /** handle an agent:* message. returns true iff the message was recognized
   *  as an agent message (so the caller knows to stop dispatching). */
  async handleMessage(ws: WebSocket, msg: any): Promise<boolean> {
    const type = msg?.type
    if (typeof type !== 'string' || !type.startsWith('agent:')) return false
    const id = msg.id
    try {
      const result = await this.dispatch(ws, type, msg)
      this.respond(ws, id, result)
    } catch (err) {
      if (err instanceof AgentSessionError) {
        this.respondError(ws, id, err.message, err.code)
      } else {
        this.respondError(ws, id, err instanceof Error ? err.message : String(err))
      }
    }
    return true
  }

  /** run once on daemon boot. idempotent — `seedFromDemoAppRegistry` no-ops
   *  when the store already has projects. */
  async seedOnBoot(): Promise<void> {
    try {
      await seedFromDemoAppRegistry()
    } catch (err) {
      process.stderr.write(
        `[sootsim-agent] seedFromDemoAppRegistry failed: ${err instanceof Error ? err.message : String(err)}\n`,
      )
    }
  }

  /** terminate every subscription and drop every socket reference. called
   *  by the bridge host during shutdown. */
  close(): void {
    if (this.sessionStatusPersistTimer) {
      clearTimeout(this.sessionStatusPersistTimer)
      this.sessionStatusPersistTimer = null
    }
    this.flushSessionStatuses()
    for (const sub of this.subscriptions.values()) {
      try {
        sub.unsubscribe()
      } catch {}
    }
    this.subscriptions.clear()
    this.sessionsBySocket.clear()
    this.allSockets.clear()
  }

  // --- dispatch ---

  private async dispatch(ws: WebSocket, type: string, msg: any): Promise<unknown> {
    switch (type) {
      case 'agent:list-projects':
        return listProjects()
      case 'agent:upsert-project':
        return upsertProject(msg.input ?? {})
      case 'agent:delete-project':
        deleteProject(String(msg.projectId))
        return { ok: true }
      case 'agent:auto-attach-for-url':
        return this.autoAttachForUrl(msg.input ?? {})
      case 'agent:list-sessions':
        return listSessions(msg.projectId ? String(msg.projectId) : undefined)
      case 'agent:start-session':
        return this.doStartSession(msg.input ?? {})
      case 'agent:send-claimed-prompt':
        return this.sendClaimedPrompt(msg)
      case 'agent:send-prompt': {
        const sessionId = String(msg.sessionId)
        const session = findSessionById(sessionId)
        if (!session) {
          throw new AgentSessionError('NO_SESSION', `no session: ${sessionId}`)
        }
        const prompt = this.normalizePromptEnvelope(msg)
        await sendAgentPrompt(sessionId, prompt)
        return this.notePromptAccepted(sessionId, prompt, session.status === 'working')
      }
      case 'agent:end-session':
        this.dropSessionFanout(String(msg.sessionId))
        await endAgentSession(String(msg.sessionId))
        const ended = findSessionById(String(msg.sessionId))
        if (ended) {
          this.broadcastSessionStatus(ended)
        }
        return { ok: true }
      case 'agent:get-transcript':
        return this.getTranscript(String(msg.sessionId))
      case 'agent:get-paths':
        return this.getPaths()
      case 'agent:subscribe-events':
        return this.subscribeSocket(ws, String(msg.sessionId))
      case 'agent:unsubscribe-events':
        return this.unsubscribeSocket(ws, String(msg.sessionId))
      default:
        throw new AgentSessionError('UNKNOWN_AGENT_MSG', `unknown agent message: ${type}`)
    }
  }

  // --- operation impls ---

  private async sendClaimedPrompt(msg: any): Promise<{
    ok: true
    routed: 'team-machine'
    sessionId: string
  }> {
    const simId = typeof msg.simId === 'string' ? msg.simId.trim() : ''
    if (!simId) {
      throw new AgentSessionError('NO_SIM', 'agent:send-claimed-prompt requires simId')
    }
    const lease = this.opts.resolveCliLease?.(simId) ?? null
    if (!lease || lease.kind !== 'cli' || lease.expiresAt <= Date.now()) {
      throw new AgentSessionError('NO_CLAIM', `sim ${simId} has no active CLI claim`)
    }
    const sessionId = teamMachineSessionIdFromCliIdentity(lease.cliIdentityKey)
    if (!sessionId) {
      throw new AgentSessionError(
        'UNSUPPORTED_CLAIM',
        `sim ${simId} is claimed by a CLI identity that is not promptable`,
      )
    }
    await sendPromptToTeamMachineSession({
      sessionId,
      prompt: this.normalizePromptEnvelope(msg),
    })
    return { ok: true, routed: 'team-machine', sessionId }
  }

  private async doStartSession(input: {
    projectId: string
    provider?: Provider
    codexBin?: string
    claudeBin?: string
    freshThread?: boolean
  }): Promise<{ session: AgentSession; wrapperPid: number }> {
    const project = findProjectById(input.projectId)
    if (!project) {
      throw new AgentSessionError('NO_PROJECT', `no project: ${input.projectId}`)
    }
    const result = await startAgentSession(input)
    this.broadcastSessionStatus(result.session)
    return result
  }

  private async autoAttachForUrl(input: {
    bundleUrl?: string
    provider?: Provider
  }): Promise<{ project: AttachedProject | null }> {
    const bundleUrl = input.bundleUrl ?? ''
    const targetPort = (() => {
      try {
        return new URL(bundleUrl).port || null
      } catch {
        return null
      }
    })()
    if (!targetPort) return { project: null }
    const excludePorts = this.opts.getExcludePorts?.() ?? defaultExcludePorts()
    const servers = await scanDevServers({ excludePorts })
    const match = servers.find((s) => String(s.port) === targetPort)
    if (!match || !match.cwd) return { project: null }
    const existing = listProjects().find((p) => p.cwd === match.cwd) ?? null
    const knownBundleUrls = Array.from(
      new Set([...(existing?.knownBundleUrls ?? []), match.bundleUrl, bundleUrl]),
    )
    const project = upsertProject({
      cwd: match.cwd,
      name: match.projectName ?? path.basename(match.cwd),
      preferredProvider: input.provider ?? existing?.preferredProvider,
      sourceRoots: existing?.sourceRoots ?? [match.cwd],
      knownBundleUrls,
      framework: existing?.framework ?? mapFrameworkToProjectFramework(match.framework),
      bundleId: match.bundleId ?? existing?.bundleId,
    })
    return { project }
  }

  private getTranscript(sessionId: string): string | { error: string; code: string } {
    const p = transcriptPath(sessionId)
    if (!fs.existsSync(p)) {
      return { error: 'transcript not found', code: 'NO_TRANSCRIPT' }
    }
    return fs.readFileSync(p, 'utf8')
  }

  private getPaths() {
    const dir = getUserDataDir()
    return {
      userDataDir: dir,
      storeFile: path.join(dir, 'attached-projects.json'),
      sessionsDir: path.join(dir, 'sessions'),
      transcriptsDir: path.join(dir, 'transcripts'),
    }
  }

  // --- subscription management ---

  private subscribeSocket(
    ws: WebSocket,
    sessionId: string,
  ): { ok: true; refCount: number } {
    let sockets = this.sessionsBySocket.get(ws)
    if (!sockets) {
      sockets = new Set()
      this.sessionsBySocket.set(ws, sockets)
    }
    if (sockets.has(sessionId)) {
      return { ok: true, refCount: this.subscriptions.get(sessionId)?.refCount ?? 1 }
    }
    sockets.add(sessionId)
    const existing = this.subscriptions.get(sessionId)
    if (existing) {
      existing.refCount++
      return { ok: true, refCount: existing.refCount }
    }
    const unsubscribe = subscribeAgentEvents(sessionId, (event) => {
      const coalesced = this.coalescePromptEcho(sessionId, event)
      if (coalesced) {
        this.applySessionEvent(sessionId, coalesced)
        this.fanOutEvent(sessionId, coalesced)
      }
      // recordTurnTelemetry mirrors what electron main used to do — keep it
      // here so cost tracking happens regardless of which client subscribed.
      if (event.type === 'turn-completed') {
        const session = findSessionById(sessionId)
        if (session) {
          try {
            recordTurnTelemetry(session.projectId, {
              usd: event.costUsd,
              ts: event.ts,
            })
          } catch (err) {
            process.stderr.write(
              `[sootsim-agent] recordTurnTelemetry failed: ${err instanceof Error ? err.message : String(err)}\n`,
            )
          }
        }
      }
    })
    this.subscriptions.set(sessionId, { unsubscribe, refCount: 1 })
    return { ok: true, refCount: 1 }
  }

  private unsubscribeSocket(
    ws: WebSocket,
    sessionId: string,
  ): { ok: true; refCount: number } {
    const sockets = this.sessionsBySocket.get(ws)
    if (!sockets || !sockets.has(sessionId)) return { ok: true, refCount: 0 }
    sockets.delete(sessionId)
    return this.decrementSubscription(sessionId)
  }

  private decrementSubscription(sessionId: string): { ok: true; refCount: number } {
    const existing = this.subscriptions.get(sessionId)
    if (!existing) return { ok: true, refCount: 0 }
    existing.refCount--
    if (existing.refCount <= 0) {
      try {
        existing.unsubscribe()
      } catch {}
      this.subscriptions.delete(sessionId)
      return { ok: true, refCount: 0 }
    }
    return { ok: true, refCount: existing.refCount }
  }

  /** end-session tears the FIFO down, so drop our reader before it
   *  disappears regardless of remaining subscriber refcount. */
  private dropSessionFanout(sessionId: string): void {
    const existing = this.subscriptions.get(sessionId)
    if (existing) {
      try {
        existing.unsubscribe()
      } catch {}
      this.subscriptions.delete(sessionId)
    }
    for (const sockets of this.sessionsBySocket.values()) {
      sockets.delete(sessionId)
    }
    this.clearPromptTracking(sessionId)
  }

  // --- wire pushes + responses ---

  private normalizePromptEnvelope(msg: any): AgentPromptEnvelope {
    if (msg?.prompt && typeof msg.prompt === 'object') {
      const prompt = msg.prompt as Record<string, unknown>
      return {
        text: String(prompt.text ?? ''),
        ...(typeof prompt.displayText === 'string'
          ? { displayText: prompt.displayText }
          : {}),
        ...(typeof prompt.inspectSummary === 'string'
          ? { inspectSummary: prompt.inspectSummary }
          : {}),
        ...(typeof prompt.inspectTrace === 'string'
          ? { inspectTrace: prompt.inspectTrace }
          : {}),
      }
    }
    return {
      text: String(msg?.text ?? ''),
      ...(typeof msg?.displayText === 'string' ? { displayText: msg.displayText } : {}),
      ...(typeof msg?.inspectSummary === 'string'
        ? { inspectSummary: msg.inspectSummary }
        : {}),
      ...(typeof msg?.inspectTrace === 'string'
        ? { inspectTrace: msg.inspectTrace }
        : {}),
    }
  }

  private notePromptAccepted(
    sessionId: string,
    prompt: AgentPromptEnvelope,
    assumeQueued: boolean,
  ): { ok: true; queued: boolean; pendingTurns: number; queueDepth: number } {
    const now = Date.now()
    const echoes = this.pendingPromptEchoes.get(sessionId) ?? []
    echoes.push({ sentAt: now })
    this.pendingPromptEchoes.set(sessionId, echoes)
    const pendingTurns =
      Math.max(this.pendingTurns.get(sessionId) ?? 0, assumeQueued ? 1 : 0) + 1
    this.pendingTurns.set(sessionId, pendingTurns)
    const promptText = prompt.displayText ?? prompt.text
    this.patchSession(sessionId, {
      lastPrompt: promptText,
      status: 'working',
      needsAttention: false,
    })
    this.fanOutEvent(sessionId, {
      type: 'prompt-received',
      text: promptText,
      ...(prompt.inspectSummary ? { inspectSummary: prompt.inspectSummary } : {}),
      ...(prompt.inspectTrace ? { inspectTrace: prompt.inspectTrace } : {}),
      ts: now,
    } as AgentEvent)
    return {
      ok: true,
      queued: pendingTurns > 1,
      pendingTurns,
      queueDepth: Math.max(0, pendingTurns - 1),
    }
  }

  private applySessionEvent(sessionId: string, event: AgentEvent): void {
    switch (event.type) {
      case 'prompt-received':
      case 'turn-started':
        this.patchSession(sessionId, {
          status: 'working',
          needsAttention: false,
        })
        return
      case 'turn-completed': {
        const pendingTurns = this.consumeSettledTurn(sessionId)
        this.patchSession(sessionId, {
          status: pendingTurns > 0 ? 'working' : 'idle',
          needsAttention: false,
          lastTurnFiles: event.filesTouched,
          currentlyEditing: undefined,
        })
        return
      }
      case 'approval-needed':
        this.patchSession(sessionId, {
          status: 'needs-attention',
          needsAttention: true,
        })
        return
      case 'error': {
        const pendingTurns = this.consumeSettledTurn(sessionId)
        this.patchSession(sessionId, {
          status: pendingTurns > 0 ? 'working' : 'needs-attention',
          needsAttention: pendingTurns <= 0,
          currentlyEditing: undefined,
        })
        return
      }
      case 'exited':
        this.clearPromptTracking(sessionId)
        this.patchSession(sessionId, {
          status: 'ended',
          needsAttention: false,
          wrapperPid: undefined,
          currentlyEditing: undefined,
        })
        return
      case 'ready':
      case 'turn-reasoning':
      case 'turn-message':
      case 'turn-plan':
      case 'tool-call':
      case 'file-edited':
      case 'file-diff-delta':
        return
    }
  }

  private patchSession(sessionId: string, patch: Partial<AgentSession>): void {
    const existing =
      this.pendingSessionStates.get(sessionId) ?? findSessionById(sessionId)
    if (!existing) return
    const updated = applySessionStatusPatch(existing, patch)
    this.pendingSessionStates.set(sessionId, updated)
    this.pendingSessionPatches.set(sessionId, {
      ...this.pendingSessionPatches.get(sessionId),
      ...patch,
    })
    this.broadcastSessionStatus(updated)
    if (this.sessionStatusPersistTimer) return
    this.sessionStatusPersistTimer = setTimeout(() => {
      this.sessionStatusPersistTimer = null
      this.flushSessionStatuses()
    }, SESSION_STATUS_PERSIST_DELAY_MS)
  }

  private flushSessionStatuses(): void {
    if (this.pendingSessionPatches.size === 0) return
    const pendingStates = this.pendingSessionStates
    const pendingPatches = this.pendingSessionPatches
    this.pendingSessionStates = new Map()
    this.pendingSessionPatches = new Map()
    try {
      updateSessionStatuses([...pendingPatches].map(([id, patch]) => ({ id, patch })))
    } catch (error) {
      for (const [id, session] of pendingStates) {
        if (!this.pendingSessionStates.has(id)) {
          this.pendingSessionStates.set(id, session)
        }
      }
      for (const [id, patch] of pendingPatches) {
        this.pendingSessionPatches.set(id, {
          ...patch,
          ...this.pendingSessionPatches.get(id),
        })
      }
      process.stderr.write(
        `[sootsim-agent] session status persistence failed: ${error instanceof Error ? error.message : String(error)}\n`,
      )
    }
  }

  private coalescePromptEcho(sessionId: string, event: AgentEvent): AgentEvent | null {
    if (event.type !== 'prompt-received') return event
    const pending = this.pendingPromptEchoes.get(sessionId)
    if (!pending || pending.length === 0) return event
    while (pending.length > 0 && Date.now() - pending[0]!.sentAt > 15_000) {
      pending.shift()
    }
    if (pending.length === 0) {
      this.pendingPromptEchoes.delete(sessionId)
      return event
    }
    pending.shift()
    if (pending.length === 0) {
      this.pendingPromptEchoes.delete(sessionId)
    } else {
      this.pendingPromptEchoes.set(sessionId, pending)
    }
    return null
  }

  private consumeSettledTurn(sessionId: string): number {
    const pendingTurns = Math.max(0, (this.pendingTurns.get(sessionId) ?? 1) - 1)
    if (pendingTurns > 0) {
      this.pendingTurns.set(sessionId, pendingTurns)
    } else {
      this.pendingTurns.delete(sessionId)
    }
    return pendingTurns
  }

  private clearPromptTracking(sessionId: string): void {
    this.pendingPromptEchoes.delete(sessionId)
    this.pendingTurns.delete(sessionId)
  }

  private fanOutEvent(sessionId: string, event: AgentEvent): void {
    const payload = JSON.stringify({ type: 'agent:event', sessionId, event })
    for (const [ws, sessions] of this.sessionsBySocket) {
      if (!sessions.has(sessionId)) continue
      if (ws.readyState !== WS_OPEN) continue
      try {
        ws.send(payload)
      } catch {}
    }
  }

  private broadcastSessionStatus(session: AgentSession): void {
    const payload = JSON.stringify({ type: 'agent:session-status', session })
    for (const ws of this.allSockets) {
      if (ws.readyState !== WS_OPEN) continue
      try {
        ws.send(payload)
      } catch {}
    }
  }

  private respond(ws: WebSocket, id: unknown, result: unknown): void {
    if (ws.readyState !== WS_OPEN) return
    try {
      ws.send(JSON.stringify({ id, result }))
    } catch {}
  }

  private respondError(ws: WebSocket, id: unknown, error: string, code?: string): void {
    if (ws.readyState !== WS_OPEN) return
    try {
      ws.send(JSON.stringify({ id, error, ...(code ? { code } : {}) }))
    } catch {}
  }
}

function mapFrameworkToProjectFramework(
  fw: DiscoveredServer['framework'],
): 'expo' | 'one' | 'rock' | 'unknown' {
  if (fw === 'expo') return 'expo'
  if (fw === 'one' || fw === 'vxrn') return 'one'
  return 'unknown'
}