import { ConfiguredInstance, DBOS } from "@dbos-inc/dbos-sdk";
import { sql } from "drizzle-orm";
import { db } from "../db";
import { party, partyMember } from "../db/schema";

type PartyMemberRow = {
	id: string;
	partyId: string;
	userId: string;
	joinedAt: Date;
	lastSeen: Date;
};

const MAX_POSITION = 50;
const SAVED_SCORE = 10;
const FOLLOWED_SCORE = 10;
const PLAYBACK_TODAY_SCORE = 5;
const PLAYBACK_WEEK_SCORE = 3;
const PLAYBACK_OLD_SCORE = 1;

type MemberScore = {
	userId: string;
	score: number;
};

type TrackEntityScore = {
	id: string;
	name: string;
	artists: { id: string; name: string }[];
	albumName?: string;
	memberScores: MemberScore[];
	memberCount: number;
};

type ArtistEntityScore = {
	id: string;
	name: string;
	memberScores: MemberScore[];
	memberCount: number;
};

type GenreEntityScore = {
	id: string;
	name: string;
	memberScores: MemberScore[];
	memberCount: number;
};

type StoryCluster = {
	memberIds: string[];
	memberCount: number;
	tracks: TrackEntityScore[];
	artists: ArtistEntityScore[];
	genres: GenreEntityScore[];
};

type PairwiseComparison = {
	userIdA: string;
	userIdB: string;
	sharedTracks: number;
	sharedArtists: number;
	sharedGenres: number;
	similarity: number;
};

type GenreDiversity = {
	userId: string;
	genreEntropy: number;
	totalGenreScore: number;
};

type GroupSummary = {
	totalMembers: number;
	mostSharedGenres: GenreEntityScore[];
	mostDiverseMember: GenreDiversity | null;
	mostAlignedPair: PairwiseComparison | null;
};

type MemberProfile = {
	userId: string;
	totalScore: number;
	genreScores: Record<string, number>;
	trackCount: number;
	artistCount: number;
};

type PartyAnalysisResult = {
	storyClusters: StoryCluster[];
	pairwise: PairwiseComparison[];
	groupSummary: GroupSummary;
	memberProfiles: MemberProfile[];
};

export class PartyAnalysisWorkflow extends ConfiguredInstance {
	@DBOS.workflow()
	async analyzeParty(partyId: string): Promise<PartyAnalysisResult> {
		const members = await this.fetchPartyMembers(partyId);
		if (members.length < 2) {
			return {
				storyClusters: [],
				pairwise: [],
				groupSummary: {
					totalMembers: members.length,
					mostSharedGenres: [],
					mostDiverseMember: null,
					mostAlignedPair: null,
				},
				memberProfiles: [],
			};
		}

		const memberInfos = members.map((m) => ({
			userId: m.userId,
			userName: "",
		}));

		const memberData = await this.fetchAllMemberData(memberInfos);
		const _memberMap = new Map(memberInfos.map((m) => [m.userId, m]));

		const trackMap = this.buildTrackEntityMap(memberData);
		const artistMap = this.buildArtistEntityMap(memberData);
		const genreMap = this.buildGenreEntityMap(memberData);

		const storyClusters = this.computeStoryClusters(
			trackMap,
			artistMap,
			genreMap,
			memberInfos.map((m) => m.userId),
		);

		const pairwise = this.computePairwise(memberData, memberInfos);

		const memberProfiles = this.buildMemberProfiles(memberData);

		const genreDiversityList = this.computeGenreDiversity(memberProfiles);

		const allGenreScores = this.aggregateAllGenreScores(memberData);
		const mostSharedGenres = this.getMostSharedGenres(
			allGenreScores,
			memberData,
			members.length,
		);

		const mostDiverseMember =
			genreDiversityList.length > 0
				? (genreDiversityList
						.sort((a, b) => b.genreEntropy - a.genreEntropy)
						.at(0) ?? null)
				: null;

		const mostAlignedPair =
			pairwise.length > 0
				? (pairwise.sort((a, b) => b.similarity - a.similarity).at(0) ?? null)
				: null;

		const groupSummary: GroupSummary = {
			totalMembers: members.length,
			mostSharedGenres,
			mostDiverseMember,
			mostAlignedPair,
		};

		await this.saveAnalysis(partyId, {
			storyClusters,
			pairwise,
			groupSummary,
			memberProfiles,
		});

		return {
			storyClusters,
			pairwise,
			groupSummary,
			memberProfiles,
		};
	}

	@DBOS.step()
	private async fetchPartyMembers(partyId: string): Promise<PartyMemberRow[]> {
		const result = await db
			.select()
			.from(partyMember)
			.where(sql`${partyMember.partyId} = ${partyId}`)
			.execute();
		return result as PartyMemberRow[];
	}

	@DBOS.step()
	private async fetchAllMemberData(
		members: { userId: string }[],
	): Promise<Map<string, MemberScores>> {
		const result = new Map<string, MemberScores>();

		for (const member of members) {
			const scores = await this.fetchMemberScores(member.userId);
			result.set(member.userId, scores);
		}

		return result;
	}

	@DBOS.step()
	private async fetchMemberScores(userId: string): Promise<MemberScores> {
		const scores: MemberScores = {
			tracks: new Map(),
			artists: new Map(),
			genreScores: new Map(),
			genreNames: new Map(),
		};

		const now = new Date();

		// Top tracks (medium_term) with position-based scoring
		const topTracks = await db.query.topTrack.findMany({
			where: {
				userId,
				timeline: "medium_term",
			},
			with: {
				track: {
					with: {
						artists: {
							with: {
								genres: true,
							},
						},
						album: true,
					},
				},
			},
			orderBy: {
				position: "asc",
			},
			limit: MAX_POSITION,
		});

		for (const t of topTracks) {
			const trackScore = MAX_POSITION - t.position + 1;
			if (!t.track) continue;
			this.addTrackScore(scores, t.track.id, trackScore, t.track);
			// Add to artist scores
			for (const artist of t.track.artists) {
				this.addArtistScore(scores, artist.id, trackScore, artist);
				// Add to genre scores
				for (const g of artist.genres) {
					this.addGenreScore(scores, g.id, trackScore, g.name);
				}
			}
		}

		// Saved tracks
		const savedTracks = await db.query.savedTrack.findMany({
			where: { userId },
			with: {
				track: {
					with: {
						artists: {
							with: {
								genres: true,
							},
						},
						album: true,
					},
				},
			},
		});

		for (const t of savedTracks) {
			if (!t.track) continue;
			this.addTrackScore(scores, t.track.id, SAVED_SCORE, t.track);
			for (const artist of t.track.artists) {
				this.addArtistScore(scores, artist.id, SAVED_SCORE / 2, artist);
				for (const g of artist.genres) {
					this.addGenreScore(scores, g.id, SAVED_SCORE / 4, g.name);
				}
			}
		}

		// Playback history
		const playbackHistory = await db.query.playbackHistory.findMany({
			where: { userId },
			with: {
				track: {
					with: {
						artists: {
							with: {
								genres: true,
							},
						},
						album: true,
					},
				},
			},
		});

		for (const h of playbackHistory) {
			if (!h.track) continue;
			const hoursSince =
				(now.getTime() - h.played_at.getTime()) / (1000 * 60 * 60);
			let playbackScore: number;
			if (hoursSince < 24) {
				playbackScore = PLAYBACK_TODAY_SCORE;
			} else if (hoursSince < 168) {
				playbackScore = PLAYBACK_WEEK_SCORE;
			} else {
				playbackScore = PLAYBACK_OLD_SCORE;
			}
			this.addTrackScore(scores, h.track.id, playbackScore, h.track);
			for (const artist of h.track.artists) {
				this.addArtistScore(scores, artist.id, playbackScore / 2, artist);
				for (const g of artist.genres) {
					this.addGenreScore(scores, g.id, playbackScore / 4, g.name);
				}
			}
		}

		// Top artists (medium_term)
		const topArtists = await db.query.topArtist.findMany({
			where: {
				userId,
				timeline: "medium_term",
			},
			with: {
				artist: {
					with: {
						genres: true,
					},
				},
			},
			orderBy: {
				position: "asc",
			},
			limit: MAX_POSITION,
		});

		for (const a of topArtists) {
			const artistScore = MAX_POSITION - a.position + 1;
			if (!a.artist) continue;
			this.addArtistScore(scores, a.artist.id, artistScore, a.artist);
			for (const g of a.artist.genres) {
				this.addGenreScore(scores, g.id, artistScore, g.name);
			}
		}

		// Followed artists
		const followedArtists = await db.query.followedArtist.findMany({
			where: { userId },
			with: {
				artist: {
					with: {
						genres: true,
					},
				},
			},
		});

		for (const fa of followedArtists) {
			if (!fa.artist) continue;
			this.addArtistScore(scores, fa.artist.id, FOLLOWED_SCORE, fa.artist);
			for (const g of fa.artist.genres) {
				this.addGenreScore(scores, g.id, FOLLOWED_SCORE, g.name);
			}
		}

		// Saved albums
		const savedAlbums = await db.query.savedAlbum.findMany({
			where: { userId },
			with: {
				album: {
					with: {
						artists: {
							with: {
								genres: true,
							},
						},
					},
				},
			},
		});

		for (const sa of savedAlbums) {
			if (!sa.album) continue;
			for (const artist of sa.album.artists) {
				this.addArtistScore(scores, artist.id, SAVED_SCORE / 2, artist);
				for (const g of artist.genres) {
					this.addGenreScore(scores, g.id, SAVED_SCORE / 4, g.name);
				}
			}
		}

		return scores;
	}

	private addTrackScore(
		scores: MemberScores,
		trackId: string,
		score: number,
		track: {
			name: string | null;
			artists: { id: string; name: string }[];
			album?: { name: string | null } | null;
		},
	): void {
		const existing = scores.tracks.get(trackId);
		if (existing) {
			existing.score += score;
		} else {
			scores.tracks.set(trackId, {
				userId: trackId,
				score,
				name: track.name ?? "",
				artists: track.artists.map((a) => ({ id: a.id, name: a.name })),
				albumName: track.album?.name ?? undefined,
			});
		}
	}

	private addArtistScore(
		scores: MemberScores,
		artistId: string,
		score: number,
		artist: { name: string },
	): void {
		const existing = scores.artists.get(artistId);
		if (existing) {
			existing.score += score;
		} else {
			scores.artists.set(artistId, {
				id: artistId,
				name: artist.name,
				score,
			});
		}
	}

	private addGenreScore(
		scores: MemberScores,
		genreId: string,
		score: number,
		genreName?: string,
	): void {
		const existing = scores.genreScores.get(genreId);
		if (existing) {
			scores.genreScores.set(genreId, existing + score);
		} else {
			scores.genreScores.set(genreId, score);
		}
		if (genreName) {
			scores.genreNames.set(genreId, genreName);
		}
	}

	private buildTrackEntityMap(
		memberData: Map<string, MemberScores>,
	): Map<string, TrackEntityScore> {
		const entityMap = new Map<string, Map<string, number>>();

		for (const [userId, data] of memberData) {
			for (const [trackId, track] of data.tracks) {
				let memberScores = entityMap.get(trackId);
				if (!memberScores) {
					memberScores = new Map();
					entityMap.set(trackId, memberScores);
				}
				memberScores.set(userId, track.score);
			}
		}

		const result = new Map<string, TrackEntityScore>();

		for (const [trackId, memberScores] of entityMap) {
			const trackInfo = Array.from(memberData.values())
				.find((data) => data.tracks.has(trackId))
				?.tracks.get(trackId) ?? {
				name: "",
				artists: [],
				albumName: undefined,
			};

			result.set(trackId, {
				id: trackId,
				name: trackInfo.name,
				artists: trackInfo.artists,
				albumName: trackInfo.albumName,
				memberScores: Array.from(memberScores.entries()).map(
					([userId, score]) => ({
						userId,
						score,
					}),
				),
				memberCount: memberScores.size,
			});
		}

		return result;
	}

	private buildArtistEntityMap(
		memberData: Map<string, MemberScores>,
	): Map<string, ArtistEntityScore> {
		const entityMap = new Map<string, Set<string>>();

		for (const [userId, data] of memberData) {
			for (const [artistId, _artistData] of data.artists) {
				let members = entityMap.get(artistId);
				if (!members) {
					members = new Set();
					entityMap.set(artistId, members);
				}
				members.add(userId);
			}
		}

		const result = new Map<string, ArtistEntityScore>();

		for (const [artistId, members] of entityMap) {
			let _totalScore = 0;
			const scoresByMember: MemberScore[] = [];

			for (const userId of members) {
				const data = memberData.get(userId);
				const artistData = data?.artists.get(artistId);
				if (artistData) {
					_totalScore += artistData.score;
					scoresByMember.push({ userId, score: artistData.score });
				}
			}

			const artistInfo = Array.from(memberData.values())
				.find((data) => data.artists.has(artistId))
				?.artists.get(artistId) ?? { name: "" };

			result.set(artistId, {
				id: artistId,
				name: artistInfo.name,
				memberScores: scoresByMember,
				memberCount: members.size,
			});
		}

		return result;
	}

	private buildGenreEntityMap(
		memberData: Map<string, MemberScores>,
	): Map<string, GenreEntityScore> {
		const entityMap = new Map<string, Map<string, number>>();
		const genreNameMap = new Map<string, string>();

		for (const [userId, data] of memberData) {
			for (const [genreId, score] of data.genreScores) {
				let memberScores = entityMap.get(genreId);
				if (!memberScores) {
					memberScores = new Map();
					entityMap.set(genreId, memberScores);
				}
				memberScores.set(userId, score);
			}
			for (const [genreId, name] of data.genreNames) {
				if (!genreNameMap.has(genreId)) {
					genreNameMap.set(genreId, name);
				}
			}
		}

		const result = new Map<string, GenreEntityScore>();

		for (const [genreId, memberScores] of entityMap) {
			result.set(genreId, {
				id: genreId,
				name: genreNameMap.get(genreId) ?? "Unknown",
				memberScores: Array.from(memberScores.entries()).map(
					([userId, score]) => ({
						userId,
						score,
					}),
				),
				memberCount: memberScores.size,
			});
		}

		return result;
	}

	private computeStoryClusters(
		trackMap: Map<string, TrackEntityScore>,
		artistMap: Map<string, ArtistEntityScore>,
		genreMap: Map<string, GenreEntityScore>,
		allMemberIds: string[],
	): StoryCluster[] {
		const clusterMap = new Map<string, StoryCluster>();

		const keyForMembers = (members: string[]) =>
			members.sort((a, b) => a.localeCompare(b)).join("|");

		const getOrCreateCluster = (memberIds: string[]): StoryCluster => {
			const key = keyForMembers([...memberIds]);
			if (!clusterMap.has(key)) {
				clusterMap.set(key, {
					memberIds,
					memberCount: memberIds.length,
					tracks: [],
					artists: [],
					genres: [],
				});
			}
			const found = clusterMap.get(key);
			if (found) return found;
			const fallback: StoryCluster = {
				memberIds,
				memberCount: memberIds.length,
				tracks: [],
				artists: [],
				genres: [],
			};
			clusterMap.set(key, fallback);
			return fallback;
		};

		// Group tracks by member subset
		for (const track of trackMap.values()) {
			const memberIds = track.memberScores.map((m) => m.userId);
			const cluster = getOrCreateCluster(memberIds);
			cluster.tracks.push(track);
		}

		// Group artists by member subset
		for (const artist of artistMap.values()) {
			const memberIds = artist.memberScores.map((m) => m.userId);
			const cluster = getOrCreateCluster(memberIds);
			cluster.artists.push(artist);
		}

		// Group genres by member subset
		for (const genre of genreMap.values()) {
			const memberIds = genre.memberScores.map((m) => m.userId);
			const cluster = getOrCreateCluster(memberIds);
			cluster.genres.push(genre);
		}

		// Sort clusters: everyone first, then by memberCount desc, then by totalScore desc
		const clusters = Array.from(clusterMap.values());
		const allMembersKey = keyForMembers([...allMemberIds]);

		clusters.sort((a, b) => {
			const aIsAll = keyForMembers([...a.memberIds]) === allMembersKey ? 1 : 0;
			const bIsAll = keyForMembers([...b.memberIds]) === allMembersKey ? 1 : 0;
			if (aIsAll !== bIsAll) return bIsAll - aIsAll;
			if (a.memberCount !== b.memberCount) return b.memberCount - a.memberCount;
			const aTotal = a.tracks.reduce(
				(s, t) => s + t.memberScores.reduce((ss, m) => ss + m.score, 0),
				0,
			);
			const bTotal = b.tracks.reduce(
				(s, t) => s + t.memberScores.reduce((ss, m) => ss + m.score, 0),
				0,
			);
			return bTotal - aTotal;
		});

		// Sort entities within each cluster
		for (const cluster of clusters) {
			cluster.tracks.sort((a, b) => {
				const aTotal = a.memberScores.reduce((s, m) => s + m.score, 0);
				const bTotal = b.memberScores.reduce((s, m) => s + m.score, 0);
				return bTotal - aTotal;
			});
			cluster.artists.sort((a, b) => {
				const aTotal = a.memberScores.reduce((s, m) => s + m.score, 0);
				const bTotal = b.memberScores.reduce((s, m) => s + m.score, 0);
				return bTotal - aTotal;
			});
			cluster.genres.sort((a, b) => {
				const aTotal = a.memberScores.reduce((s, m) => s + m.score, 0);
				const bTotal = b.memberScores.reduce((s, m) => s + m.score, 0);
				return bTotal - aTotal;
			});
		}

		return clusters;
	}

	private computePairwise(
		memberData: Map<string, MemberScores>,
		members: { userId: string }[],
	): PairwiseComparison[] {
		const pairwise: PairwiseComparison[] = [];

		for (const [i, memberA] of members.entries()) {
			for (const [j, memberB] of members.entries()) {
				if (j <= i) continue;

				const dataA = memberData.get(memberA.userId);
				const dataB = memberData.get(memberB.userId);
				if (!dataA || !dataB) continue;

				// Shared tracks
				let sharedTracks = 0;
				let intersectionScore = 0;
				let unionScore = 0;

				const allTrackIds = new Set([
					...dataA.tracks.keys(),
					...dataB.tracks.keys(),
				]);
				for (const trackId of allTrackIds) {
					const scoreA = dataA.tracks.get(trackId)?.score ?? 0;
					const scoreB = dataB.tracks.get(trackId)?.score ?? 0;
					if (scoreA > 0 && scoreB > 0) {
						sharedTracks++;
						intersectionScore += Math.min(scoreA, scoreB);
					}
					unionScore += Math.max(scoreA, scoreB);
				}

				// Shared artists
				let sharedArtists = 0;
				const allArtistIds = new Set([
					...dataA.artists.keys(),
					...dataB.artists.keys(),
				]);
				for (const artistId of allArtistIds) {
					const scoreA = dataA.artists.get(artistId)?.score ?? 0;
					const scoreB = dataB.artists.get(artistId)?.score ?? 0;
					if (scoreA > 0 && scoreB > 0) {
						sharedArtists++;
						intersectionScore += Math.min(scoreA, scoreB);
					}
					unionScore += Math.max(scoreA, scoreB);
				}

				// Shared genres
				let sharedGenres = 0;
				const allGenreIds = new Set([
					...dataA.genreScores.keys(),
					...dataB.genreScores.keys(),
				]);
				for (const genreId of allGenreIds) {
					const scoreA = dataA.genreScores.get(genreId) ?? 0;
					const scoreB = dataB.genreScores.get(genreId) ?? 0;
					if (scoreA > 0 && scoreB > 0) {
						sharedGenres++;
						intersectionScore += Math.min(scoreA, scoreB);
					}
					unionScore += Math.max(scoreA, scoreB);
				}

				const similarity = unionScore > 0 ? intersectionScore / unionScore : 0;

				pairwise.push({
					userIdA: memberA.userId,
					userIdB: memberB.userId,
					sharedTracks,
					sharedArtists,
					sharedGenres,
					similarity,
				});
			}
		}

		return pairwise;
	}

	private buildMemberProfiles(
		memberData: Map<string, MemberScores>,
	): MemberProfile[] {
		const profiles: MemberProfile[] = [];

		for (const [userId, data] of memberData) {
			const totalScore =
				[...data.tracks.values()].reduce((s, t) => s + t.score, 0) +
				[...data.artists.values()].reduce((s, a) => s + a.score, 0);

			profiles.push({
				userId,
				totalScore,
				genreScores: Object.fromEntries(data.genreScores),
				trackCount: data.tracks.size,
				artistCount: data.artists.size,
			});
		}

		return profiles;
	}

	private computeGenreDiversity(
		memberProfiles: MemberProfile[],
	): GenreDiversity[] {
		const diversity: GenreDiversity[] = [];

		for (const profile of memberProfiles) {
			const genreScores = profile.genreScores;
			const totalScore = profile.totalScore;

			if (totalScore === 0) {
				diversity.push({
					userId: profile.userId,
					genreEntropy: 0,
					totalGenreScore: 0,
				});
				continue;
			}

			let entropy = 0;
			for (const score of Object.values(genreScores)) {
				const p = score / totalScore;
				if (p > 0) {
					entropy -= p * Math.log(p);
				}
			}

			diversity.push({
				userId: profile.userId,
				genreEntropy: entropy,
				totalGenreScore: totalScore,
			});
		}

		return diversity;
	}

	private aggregateAllGenreScores(
		memberData: Map<string, MemberScores>,
	): Map<string, Map<string, number>> {
		const genreMap = new Map<string, Map<string, number>>();

		for (const [userId, data] of memberData) {
			for (const [genreId, score] of data.genreScores) {
				let memberScores = genreMap.get(genreId);
				if (!memberScores) {
					memberScores = new Map();
					genreMap.set(genreId, memberScores);
				}
				memberScores.set(userId, score);
			}
		}

		return genreMap;
	}

	private getMostSharedGenres(
		genreMap: Map<string, Map<string, number>>,
		memberData: Map<string, MemberScores>,
		_totalMembers: number,
	): GenreEntityScore[] {
		const genreNameMap = new Map<string, string>();
		for (const [, data] of memberData) {
			for (const [genreId, name] of data.genreNames) {
				if (!genreNameMap.has(genreId)) {
					genreNameMap.set(genreId, name);
				}
			}
		}

		const genres: GenreEntityScore[] = [];

		for (const [genreId, memberScores] of genreMap) {
			genres.push({
				id: genreId,
				name: genreNameMap.get(genreId) ?? "Unknown",
				memberScores: Array.from(memberScores.entries()).map(
					([userId, score]) => ({
						userId,
						score,
					}),
				),
				memberCount: memberScores.size,
			});
		}

		genres.sort(
			(a, b) =>
				b.memberCount - a.memberCount ||
				b.memberScores.reduce((total, member) => total + member.score, 0) -
					a.memberScores.reduce((total, member) => total + member.score, 0),
		);

		// Return top genres that are shared by at least 2 members
		return genres.filter((g) => g.memberCount >= 2).slice(0, 10);
	}

	@DBOS.step()
	private async saveAnalysis(
		partyId: string,
		analysis: PartyAnalysisResult,
	): Promise<void> {
		await db
			.update(party)
			.set({
				analysisData: analysis,
				lastUpdated: new Date(),
			})
			.where(sql`${party.id} = ${partyId}`);
	}
}

interface MemberScores {
	tracks: Map<
		string,
		{
			userId: string;
			score: number;
			name: string;
			artists: { id: string; name: string }[];
			albumName?: string;
		}
	>;
	artists: Map<string, { id: string; name: string; score: number }>;
	genreScores: Map<string, number>;
	genreNames: Map<string, string>;
}

export const partyAnalysisWorkflow = new PartyAnalysisWorkflow(
	"party-analysis",
);