import type { CancellationToken } from 'vscode';
import type { Source } from '../../../../constants.telemetry.js';
import type { GitRepositoryService } from '../../../../git/gitRepositoryService.js';
import { ComposeToolsIntegration } from '../../../../plus/coretools/compose/integration.js';
import type {
	ComposeApplyPlan,
	ComposeHunk,
	ComposePlan,
	ComposePlanResult,
	ComposeProgressEvent,
	ComposeSource,
	ComposeTarget,
	SigningConfig,
} from '../../../../plus/coretools/compose/types.js';
import { applyComposePlan, cancellationTokenToSignal, composePlan } from '../../../../plus/coretools/compose/utils.js';
import type { ComposerCommit, ComposerHunk } from '../protocol.js';

export interface GeneratePlanInput {
	svc: GitRepositoryService;
	source: ComposerSource;
	target?: ComposeTarget;
	customInstructions?: string;
	context?: { issueUrl?: string; pullRequestUrl?: string };
	cancellation?: CancellationToken;
	onProgress?: (event: ComposeProgressEvent) => void;
	telemetrySource: Source;
	/**
	 * When true, skips the large-prompt warning modal. Set after the first
	 * confirmed compose in a session so subsequent calls don't re-prompt.
	 */
	suppressLargePromptWarning?: boolean;
}

export interface GeneratePlanResult {
	/** The plan snapshot to cache — pass back to {@link ComposerComposeIntegration.applyPlan}. */
	cacheKey: string;
	/** Hunks the library collected, converted to GitLens's ComposerHunk shape. */
	hunks: ComposerHunk[];
	/** Commits generated by the AI, as GitLens's ComposerCommit shape. */
	commits: ComposerCommit[];
	/** Diagnostic stats for telemetry. */
	diffStats: { fileCount: number; hunkCount: number; addedLines: number; removedLines: number };
	/** Aggregated AI usage across all library phases. */
	usage: { inputTokens: number; outputTokens: number };
}

export interface ApplyPlanInput {
	svc: GitRepositoryService;
	/** Key from the previous {@link ComposerComposeIntegration.generatePlan} call. */
	cacheKey: string;
	/** User-edited commits from the webview. */
	commits: ComposerCommit[];
	signing?: SigningConfig;
	cancellation?: CancellationToken;
	onProgress?: (event: ComposeProgressEvent) => void;
	telemetrySource: Source;
}

export interface ApplyPlanResult {
	commitShas?: Record<string, string>;
	undoId?: string;
}

export type ComposerSource =
	| { type: 'workdir'; stagedOnly?: boolean; noUntracked?: boolean }
	| { type: 'branch'; name: string; mergeTarget?: string }
	| { type: 'commit-range'; branch: string; from: string; to: string };

/**
 * Compose-tools integration specialized for the legacy Composer webview.
 *
 * Wraps the library in shapes the composer protocol speaks (`ComposerHunk`,
 * `ComposerCommit`, `ComposerSource`). The two-phase protocol: {@link generatePlan}
 * returns a `cacheKey` that must be passed to {@link applyPlan} later. The base
 * class owns the cache keyed by that string.
 */
export class ComposerComposeIntegration extends ComposeToolsIntegration {
	async generatePlan(input: GeneratePlanInput): Promise<GeneratePlanResult> {
		const git = this.createGitPort(input.svc);
		const model = this.createAiModelPort(input.telemetrySource);
		const { signal, dispose: disposeSignal } = cancellationTokenToSignal(input.cancellation);
		const onBeforePrompt = this.buildLargePromptGate(input.suppressLargePromptWarning ?? false);

		const source = this.toLibrarySource(input.source);
		const target = input.target;

		try {
			const result: ComposePlanResult = await composePlan({
				git: git,
				model: model,
				source: source,
				target: target,
				instructions: input.customInstructions,
				context: input.context,
				onProgress: input.onProgress,
				cancellation: signal,
				onBeforePrompt: onBeforePrompt,
			});

			const cacheKey = this.createCacheKey(input.svc.path);
			this._cache.set(cacheKey, {
				plan: result.plan,
				snapshot: result.snapshot,
				source: source,
				sourceHunks: result.source.hunks,
			});

			return {
				cacheKey: cacheKey,
				hunks: result.source.hunks.map(h => this.toComposerHunk(h)),
				commits: result.plan.allOrderedCommits.map(c => this.toComposerCommit(c)),
				diffStats: {
					fileCount: new Set(result.source.hunks.map(h => h.fileName)).size,
					hunkCount: result.source.hunks.length,
					addedLines: result.source.hunks.reduce((sum, h) => sum + h.additions, 0),
					removedLines: result.source.hunks.reduce((sum, h) => sum + h.deletions, 0),
				},
				usage: {
					inputTokens: result.usage.inputTokens,
					outputTokens: result.usage.outputTokens,
				},
			};
		} finally {
			disposeSignal();
		}
	}

	async applyPlan(input: ApplyPlanInput): Promise<ApplyPlanResult> {
		const cached = this._cache.get(input.cacheKey);
		if (!cached) {
			throw new Error(`No cached compose plan for key '${input.cacheKey}'. Call generatePlan() first.`);
		}

		const git = this.createGitPort(input.svc);

		const userEditedPlan = this.reconstructPlan(cached.plan, input.commits);

		const applyPlanInput: ComposeApplyPlan = {
			plan: userEditedPlan,
			source: cached.source,
			snapshot: cached.snapshot,
		};

		try {
			const result = await applyComposePlan({
				git: git,
				applyPlan: applyPlanInput,
				onProgress: input.onProgress,
				signing: input.signing,
				// For recompose flows: pass 'plurality' so commits with hunks from multiple
				// authors get correct attribution + Co-authored-by trailers. For workdir
				// source, all hunks lack authors so this is a no-op.
				authorAttribution: 'plurality',
			});

			return {
				commitShas: result.commitShas,
				undoId: result.undoId,
			};
		} finally {
			// One-shot: drop the cached plan once apply runs. Another compose click
			// will produce a fresh cacheKey.
			this._cache.delete(input.cacheKey);
		}
	}

	private toLibrarySource(source: ComposerSource): ComposeSource {
		if (source.type === 'workdir') {
			return { type: 'workdir', stagedOnly: source.stagedOnly, noUntracked: source.noUntracked };
		}
		if (source.type === 'commit-range') {
			return { type: 'commit-range', branch: source.branch, from: source.from, to: source.to };
		}
		return { type: 'branch', name: source.name, mergeTarget: source.mergeTarget };
	}

	private toComposerHunk(h: ComposeHunk): ComposerHunk {
		// compose-tools no longer attaches per-hunk source-layer info (combined-diff
		// collection drops the distinction). Use the 'unknown' sentinel on the GitLens
		// internal type so downstream categorizers can tell library-converted hunks
		// apart from the webview composer's own staged/unstaged-tagged hunks.
		//
		// GitCommitIdentityShape requires a `date: Date`; the library's identity
		// carries an optional date (workdir hunks have none). Synthesize `new Date()`
		// as a neutral stand-in so the shape lines up — commit creation only uses
		// the date when it came from a real commit-range hunk.
		//
		// Indexing: the library emits 0-based hunk indices; GitLens's `ComposerHunk`
		// protocol is 1-based (see `protocol.ts`). Shift by +1 here and in
		// `toComposerCommit`; shift back by -1 in `reconstructPlan`.
		return {
			index: h.index + 1,
			fileName: h.fileName,
			diffHeader: h.diffHeader,
			hunkHeader: h.hunkHeader,
			content: h.content,
			additions: h.additions,
			deletions: h.deletions,
			source: 'unknown',
			isRename: h.isRename,
			originalFileName: h.originalFileName,
			author: h.author
				? { name: h.author.name, email: h.author.email, date: h.author.date ?? new Date() }
				: undefined,
			coAuthors: h.coAuthors?.map(ca => ({
				name: ca.name,
				email: ca.email,
				date: ca.date ?? new Date(),
			})),
		};
	}

	private toComposerCommit(c: {
		id: string;
		message: string;
		explanation: string;
		hunkIndices: number[];
	}): ComposerCommit {
		return {
			id: c.id,
			message: { content: c.message, isGenerated: true },
			aiExplanation: c.explanation,
			// Library indices are 0-based; GitLens's protocol is 1-based. Keep the
			// same +1 shift used in `toComposerHunk` so commit references line up
			// with the hunks that back them.
			hunkIndices: c.hunkIndices.map(i => i + 1),
		};
	}

	/**
	 * Rebuild a ComposePlan from the AI-generated plan + user edits.
	 *
	 * The user may have reassigned hunks between commits, reordered commits, or
	 * edited messages. We preserve the original branch-group scaffolding but
	 * replace `allOrderedCommits` and each branch's `commits` with the edited
	 * versions. Hunk indices still reference the library's original sourceHunks —
	 * the user didn't modify which hunks exist, only which commit they belong to.
	 *
	 * Shift indices back by -1 to undo the +1 applied in `toComposerHunk` /
	 * `toComposerCommit`. The library works in 0-based space internally.
	 */
	private reconstructPlan(originalPlan: ComposePlan, userCommits: ComposerCommit[]): ComposePlan {
		const newCommits = userCommits.map(uc => ({
			id: uc.id,
			message: uc.message.content,
			explanation: uc.aiExplanation ?? '',
			hunkIndices: uc.hunkIndices.map(i => i - 1),
		}));

		// For target: 'head' (single branch), put all commits in one branch.
		// For multi-branch targets the library's original branch structure is
		// kept; we slot the user's commits into the first branch.
		const rebuiltBranches = originalPlan.branches.map((bp, idx) => ({
			branchGroup: {
				...bp.branchGroup,
				commits: idx === 0 ? newCommits : bp.branchGroup.commits,
			},
			orderedCommitIds: idx === 0 ? newCommits.map(c => c.id) : bp.orderedCommitIds,
			partitioning: bp.partitioning,
		}));

		return {
			analysis: originalPlan.analysis,
			grouping: {
				branches: rebuiltBranches.map(bp => bp.branchGroup),
			},
			ordering: {
				...originalPlan.ordering,
				branches: rebuiltBranches.map(bp => ({
					branchId: bp.branchGroup.id,
					orderedCommitIds: bp.orderedCommitIds,
				})),
			},
			branches: rebuiltBranches,
			allOrderedCommits: newCommits,
		};
	}
}
