import { type DraftStore, makeDraftStore } from '../../../authoring/index.ts'
import { type CaptureSession, startCapture } from '../../../cdp/capture.ts'
import { defineTool, type RegisteredTool } from '../../server.ts'
import type { AttachState } from './attach.ts'

/**
 * One capture == a session (network requests) + the tab CDP client it runs on +
 * a draft store. Bundling the tab client is what lets the shared authoring
 * tools treat a `captureId` as a full session (drive the browser AND read
 * capture AND hold drafts) behind one id.
 */
export interface CaptureEntry {
	session: CaptureSession
	tabClient: Awaited<ReturnType<AttachState['handles']['attachToTab']>>
	drafts: DraftStore
}

export interface CaptureStore {
	sessions: Map<string, CaptureEntry>
	activeTab?: string
}

export async function resetCaptures(captures: CaptureStore) {
	for(const entry of captures.sessions.values()) {
		await entry.session.stop().catch(() => {})
	}

	captures.sessions.clear()
	delete captures.activeTab
}

export function captureTools(
	attachRef: { current?: AttachState },
	captures: CaptureStore,
): RegisteredTool[] {
	const requireAttached = (): AttachState => {
		if(!attachRef.current) {
			throw new Error('Not attached. Call attach_browser first.')
		}

		return attachRef.current
	}

	const startTool = defineTool<{ tabId: string }>(
		{
			name: 'start_capture',
			description:
				'Begin recording network requests on the specified tab. '
				+ 'Returns a captureId — the session handle the authoring tools '
				+ '(navigate, propose_provider, run_proof, …) all take.',
			inputSchema: {
				type: 'object',
				properties: { tabId: { type: 'string' } },
				required: ['tabId'],
			},
		},
		async({ tabId }) => {
			const attached = requireAttached()
			const tabClient = await attached.handles.attachToTab(tabId)
			const session = await startCapture(tabClient)
			captures.sessions.set(session.captureId, {
				session, tabClient, drafts: makeDraftStore(),
			})
			captures.activeTab = tabId
			return { captureId: session.captureId }
		},
	)

	const stopTool = defineTool<{ captureId: string }>(
		{
			name: 'stop_capture',
			description:
				'Stop recording on the given capture. '
				+ 'Captured data remains queryable until reset_session.',
			inputSchema: {
				type: 'object',
				properties: { captureId: { type: 'string' } },
				required: ['captureId'],
			},
		},
		async({ captureId }) => {
			const entry = captures.sessions.get(captureId)
			if(!entry) {
				throw new Error(`no capture ${captureId}`)
			}

			await entry.session.stop()
			return { stopped: true, requestCount: entry.session.requests.size }
		},
	)

	const getTool = defineTool<{ captureId: string, requestId: string }>(
		{
			name: 'get_request',
			description:
				'Return the full headers and bodies of one captured request. Use it '
				+ 'to inspect a candidate that find_requests_containing ranked. '
				+ 'Unlike every other tool here, the headers are RAW — cookie and '
				+ 'authorization values are visible. Read them; never repeat them to '
				+ 'the developer or write them to a file.',
			inputSchema: {
				type: 'object',
				properties: {
					captureId: { type: 'string' },
					requestId: { type: 'string' },
				},
				required: ['captureId', 'requestId'],
			},
		},
		async({ captureId, requestId }) => {
			const entry = captures.sessions.get(captureId)
			if(!entry) {
				throw new Error(`no capture ${captureId}`)
			}

			const r = entry.session.requests.get(requestId)
			if(!r) {
				throw new Error(`no request ${requestId} in capture ${captureId}`)
			}

			const body = entry.session.bodies.get(requestId)
			if(!body.available) {
				return { ...r, bodyAvailable: false, reason: body.reason }
			}

			return { ...r, responseBody: body.body, bodyAvailable: true }
		},
	)

	return [startTool, stopTool, getTool]
}