import type { AgentState } from "@mariozechner/pi-agent-core";
import { Store } from "../store.js";
import type { SessionData, SessionMetadata, StoreConfig } from "../types.js";

/**
 * Store for chat sessions (data and metadata).
 * Uses two object stores: sessions (full data) and sessions-metadata (lightweight).
 */
export class SessionsStore extends Store {
	getConfig(): StoreConfig {
		return {
			name: "sessions",
			keyPath: "id",
			indices: [{ name: "lastModified", keyPath: "lastModified" }],
		};
	}

	/**
	 * Additional config for sessions-metadata store.
	 * Must be included when creating the backend.
	 */
	static getMetadataConfig(): StoreConfig {
		return {
			name: "sessions-metadata",
			keyPath: "id",
			indices: [{ name: "lastModified", keyPath: "lastModified" }],
		};
	}

	async save(data: SessionData, metadata: SessionMetadata): Promise<void> {
		await this.getBackend().transaction(["sessions", "sessions-metadata"], "readwrite", async (tx) => {
			await tx.set("sessions", data.id, data);
			await tx.set("sessions-metadata", metadata.id, metadata);
		});
	}

	async get(id: string): Promise<SessionData | null> {
		return this.getBackend().get("sessions", id);
	}

	async getMetadata(id: string): Promise<SessionMetadata | null> {
		return this.getBackend().get("sessions-metadata", id);
	}

	async getAllMetadata(): Promise<SessionMetadata[]> {
		// Use the lastModified index to get sessions sorted by most recent first
		return this.getBackend().getAllFromIndex<SessionMetadata>("sessions-metadata", "lastModified", "desc");
	}

	async delete(id: string): Promise<void> {
		await this.getBackend().transaction(["sessions", "sessions-metadata"], "readwrite", async (tx) => {
			await tx.delete("sessions", id);
			await tx.delete("sessions-metadata", id);
		});
	}

	// Alias for backward compatibility
	async deleteSession(id: string): Promise<void> {
		return this.delete(id);
	}

	async updateTitle(id: string, title: string): Promise<void> {
		const metadata = await this.getMetadata(id);
		if (metadata) {
			metadata.title = title;
			await this.getBackend().set("sessions-metadata", id, metadata);
		}

		// Also update in full session data
		const data = await this.get(id);
		if (data) {
			data.title = title;
			await this.getBackend().set("sessions", id, data);
		}
	}

	async getQuotaInfo(): Promise<{ usage: number; quota: number; percent: number }> {
		return this.getBackend().getQuotaInfo();
	}

	async requestPersistence(): Promise<boolean> {
		return this.getBackend().requestPersistence();
	}

	// Alias methods for backward compatibility
	async saveSession(
		id: string,
		state: AgentState,
		metadata: SessionMetadata | undefined,
		title?: string,
	): Promise<void> {
		// If metadata is provided, use it; otherwise create it from state
		const meta: SessionMetadata = metadata || {
			id,
			title: title || "",
			createdAt: new Date().toISOString(),
			lastModified: new Date().toISOString(),
			messageCount: state.messages?.length || 0,
			usage: {
				input: 0,
				output: 0,
				cacheRead: 0,
				cacheWrite: 0,
				totalTokens: 0,
				cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
			},
			thinkingLevel: state.thinkingLevel || "off",
			preview: "",
		};

		const data: SessionData = {
			id,
			title: title || meta.title,
			model: state.model,
			thinkingLevel: state.thinkingLevel,
			messages: state.messages || [],
			createdAt: meta.createdAt,
			lastModified: new Date().toISOString(),
		};

		await this.save(data, meta);
	}

	async loadSession(id: string): Promise<SessionData | null> {
		return this.get(id);
	}

	async getLatestSessionId(): Promise<string | null> {
		const allMetadata = await this.getAllMetadata();
		if (allMetadata.length === 0) return null;

		// Sort by lastModified descending
		allMetadata.sort((a, b) => b.lastModified.localeCompare(a.lastModified));
		return allMetadata[0].id;
	}
}
