import * as React from "react";
import { api } from "../api/client";
import type { TokenUsageRecord } from "../api/client";

type SortKey =
	| "timestamp"
	| "story_id"
	| "agent_name"
	| "model"
	| "total_cost_usd";
type SortDir = "asc" | "desc";

function formatCost(usd: number): string {
	if (usd === 0) return "$0.00";
	if (usd < 0.001) return `$${usd.toFixed(6)}`;
	if (usd < 0.01) return `$${usd.toFixed(4)}`;
	return `$${usd.toFixed(3)}`;
}

function formatTokens(n: number): string {
	if (n >= 1_000_000) return `${(n / 1_000_000).toFixed(1)}M`;
	if (n >= 1_000) return `${(n / 1_000).toFixed(1)}K`;
	return String(n);
}

function formatTimestamp(iso: string): string {
	const d = new Date(iso);
	const year = d.getFullYear();
	const month = String(d.getMonth() + 1).padStart(2, "0");
	const day = String(d.getDate()).padStart(2, "0");
	const h = String(d.getHours()).padStart(2, "0");
	const m = String(d.getMinutes()).padStart(2, "0");
	return `${year}-${month}-${day} ${h}:${m}`;
}

/** Infer an agent type from the agent name. */
function agentType(agentName: string): string {
	const lower = agentName.toLowerCase();
	if (lower.startsWith("coder")) return "coder";
	if (lower.startsWith("qa")) return "qa";
	if (lower.startsWith("mergemaster") || lower.startsWith("merge"))
		return "mergemaster";
	return "other";
}

interface SortHeaderProps {
	label: string;
	sortKey: SortKey;
	current: SortKey;
	dir: SortDir;
	onSort: (key: SortKey) => void;
	align?: "left" | "right";
}

function SortHeader({
	label,
	sortKey,
	current,
	dir,
	onSort,
	align = "left",
}: SortHeaderProps) {
	const active = current === sortKey;
	return (
		<th
			style={{
				padding: "8px 12px",
				textAlign: align,
				cursor: "pointer",
				userSelect: "none",
				borderBottom: "1px solid #333",
				color: active ? "#ececec" : "#aaa",
				fontWeight: active ? "700" : "500",
				whiteSpace: "nowrap",
				fontSize: "0.8em",
				letterSpacing: "0.05em",
				textTransform: "uppercase",
			}}
			onClick={() => onSort(sortKey)}
		>
			{label}
			{active ? (dir === "asc" ? " ↑" : " ↓") : ""}
		</th>
	);
}

interface TokenUsagePageProps {
	projectPath: string;
}

export function TokenUsagePage({ projectPath }: TokenUsagePageProps) {
	const [records, setRecords] = React.useState<TokenUsageRecord[]>([]);
	const [loading, setLoading] = React.useState(true);
	const [error, setError] = React.useState<string | null>(null);
	const [sortKey, setSortKey] = React.useState<SortKey>("timestamp");
	const [sortDir, setSortDir] = React.useState<SortDir>("desc");

	React.useEffect(() => {
		setLoading(true);
		setError(null);
		api
			.getAllTokenUsage()
			.then((resp) => setRecords(resp.records))
			.catch((e) =>
				setError(e instanceof Error ? e.message : "Failed to load token usage"),
			)
			.finally(() => setLoading(false));
	}, [projectPath]);

	function handleSort(key: SortKey) {
		if (key === sortKey) {
			setSortDir((d) => (d === "asc" ? "desc" : "asc"));
		} else {
			setSortKey(key);
			setSortDir(key === "timestamp" ? "desc" : "asc");
		}
	}

	const sorted = React.useMemo(() => {
		return [...records].sort((a, b) => {
			let cmp = 0;
			switch (sortKey) {
				case "timestamp":
					cmp = a.timestamp.localeCompare(b.timestamp);
					break;
				case "story_id":
					cmp = a.story_id.localeCompare(b.story_id);
					break;
				case "agent_name":
					cmp = a.agent_name.localeCompare(b.agent_name);
					break;
				case "model":
					cmp = (a.model ?? "").localeCompare(b.model ?? "");
					break;
				case "total_cost_usd":
					cmp = a.total_cost_usd - b.total_cost_usd;
					break;
			}
			return sortDir === "asc" ? cmp : -cmp;
		});
	}, [records, sortKey, sortDir]);

	// Compute summary totals
	const totalCost = records.reduce((s, r) => s + r.total_cost_usd, 0);

	const byAgentType = React.useMemo(() => {
		const map: Record<string, number> = {};
		for (const r of records) {
			const t = agentType(r.agent_name);
			map[t] = (map[t] ?? 0) + r.total_cost_usd;
		}
		return map;
	}, [records]);

	const byModel = React.useMemo(() => {
		const map: Record<string, number> = {};
		for (const r of records) {
			const m = r.model ?? "unknown";
			map[m] = (map[m] ?? 0) + r.total_cost_usd;
		}
		return map;
	}, [records]);

	const cellStyle: React.CSSProperties = {
		padding: "7px 12px",
		borderBottom: "1px solid #222",
		fontSize: "0.85em",
		color: "#ccc",
		whiteSpace: "nowrap",
	};

	return (
		<div
			style={{
				height: "100%",
				overflowY: "auto",
				background: "#111",
				padding: "24px",
				fontFamily: "monospace",
			}}
		>
			<h2
				style={{
					color: "#ececec",
					margin: "0 0 20px",
					fontSize: "1.1em",
					fontWeight: "700",
					letterSpacing: "0.04em",
				}}
			>
				Token Usage
			</h2>

			{/* Summary totals */}
			<div
				style={{
					display: "flex",
					gap: "16px",
					flexWrap: "wrap",
					marginBottom: "24px",
				}}
			>
				<SummaryCard
					label="Total Cost"
					value={formatCost(totalCost)}
					highlight
				/>
				{Object.entries(byAgentType)
					.sort(([a], [b]) => a.localeCompare(b))
					.map(([type, cost]) => (
						<SummaryCard
							key={type}
							label={`${type.charAt(0).toUpperCase()}${type.slice(1)}`}
							value={formatCost(cost)}
						/>
					))}
				{Object.entries(byModel)
					.sort(([, a], [, b]) => b - a)
					.map(([model, cost]) => (
						<SummaryCard key={model} label={model} value={formatCost(cost)} />
					))}
			</div>

			{loading && (
				<p style={{ color: "#555", fontSize: "0.9em" }}>Loading...</p>
			)}
			{error && <p style={{ color: "#e05c5c", fontSize: "0.9em" }}>{error}</p>}

			{!loading && !error && records.length === 0 && (
				<p style={{ color: "#555", fontSize: "0.9em" }}>
					No token usage records found.
				</p>
			)}

			{!loading && !error && records.length > 0 && (
				<div style={{ overflowX: "auto" }}>
					<table
						style={{
							width: "100%",
							borderCollapse: "collapse",
							fontSize: "0.9em",
						}}
					>
						<thead>
							<tr style={{ background: "#1a1a1a" }}>
								<SortHeader
									label="Date"
									sortKey="timestamp"
									current={sortKey}
									dir={sortDir}
									onSort={handleSort}
								/>
								<SortHeader
									label="Story"
									sortKey="story_id"
									current={sortKey}
									dir={sortDir}
									onSort={handleSort}
								/>
								<SortHeader
									label="Agent"
									sortKey="agent_name"
									current={sortKey}
									dir={sortDir}
									onSort={handleSort}
								/>
								<SortHeader
									label="Model"
									sortKey="model"
									current={sortKey}
									dir={sortDir}
									onSort={handleSort}
								/>
								<th
									style={{
										...cellStyle,
										borderBottom: "1px solid #333",
										textAlign: "right",
										color: "#aaa",
										fontSize: "0.8em",
										letterSpacing: "0.05em",
										textTransform: "uppercase",
										fontWeight: "500",
									}}
								>
									Input
								</th>
								<th
									style={{
										...cellStyle,
										borderBottom: "1px solid #333",
										textAlign: "right",
										color: "#aaa",
										fontSize: "0.8em",
										letterSpacing: "0.05em",
										textTransform: "uppercase",
										fontWeight: "500",
									}}
								>
									Cache+
								</th>
								<th
									style={{
										...cellStyle,
										borderBottom: "1px solid #333",
										textAlign: "right",
										color: "#aaa",
										fontSize: "0.8em",
										letterSpacing: "0.05em",
										textTransform: "uppercase",
										fontWeight: "500",
									}}
								>
									Cache↩
								</th>
								<th
									style={{
										...cellStyle,
										borderBottom: "1px solid #333",
										textAlign: "right",
										color: "#aaa",
										fontSize: "0.8em",
										letterSpacing: "0.05em",
										textTransform: "uppercase",
										fontWeight: "500",
									}}
								>
									Output
								</th>
								<SortHeader
									label="Cost"
									sortKey="total_cost_usd"
									current={sortKey}
									dir={sortDir}
									onSort={handleSort}
									align="right"
								/>
							</tr>
						</thead>
						<tbody>
							{sorted.map((r, i) => (
								<tr
									key={`${r.story_id}-${r.agent_name}-${r.timestamp}`}
									style={{ background: i % 2 === 0 ? "#111" : "#161616" }}
								>
									<td style={cellStyle}>{formatTimestamp(r.timestamp)}</td>
									<td
										style={{
											...cellStyle,
											color: "#8b9cf7",
											maxWidth: "220px",
											overflow: "hidden",
											textOverflow: "ellipsis",
										}}
									>
										{r.story_id}
									</td>
									<td style={{ ...cellStyle, color: "#7ec8a4" }}>
										{r.agent_name}
									</td>
									<td style={{ ...cellStyle, color: "#c9a96e" }}>
										{r.model ?? "—"}
									</td>
									<td style={{ ...cellStyle, textAlign: "right" }}>
										{formatTokens(r.input_tokens)}
									</td>
									<td style={{ ...cellStyle, textAlign: "right" }}>
										{formatTokens(r.cache_creation_input_tokens)}
									</td>
									<td style={{ ...cellStyle, textAlign: "right" }}>
										{formatTokens(r.cache_read_input_tokens)}
									</td>
									<td style={{ ...cellStyle, textAlign: "right" }}>
										{formatTokens(r.output_tokens)}
									</td>
									<td
										style={{
											...cellStyle,
											textAlign: "right",
											color: "#e08c5c",
											fontWeight: "600",
										}}
									>
										{formatCost(r.total_cost_usd)}
									</td>
								</tr>
							))}
						</tbody>
					</table>
				</div>
			)}
		</div>
	);
}

function SummaryCard({
	label,
	value,
	highlight = false,
}: { label: string; value: string; highlight?: boolean }) {
	return (
		<div
			style={{
				background: highlight ? "#1e1e2e" : "#1a1a1a",
				border: `1px solid ${highlight ? "#3a3a5a" : "#2a2a2a"}`,
				borderRadius: "8px",
				padding: "12px 16px",
				minWidth: "120px",
			}}
		>
			<div
				style={{
					fontSize: "0.7em",
					color: "#666",
					textTransform: "uppercase",
					letterSpacing: "0.07em",
					marginBottom: "4px",
				}}
			>
				{label}
			</div>
			<div
				style={{
					fontSize: "1.1em",
					fontWeight: "700",
					color: highlight ? "#c9a96e" : "#ececec",
					fontFamily: "monospace",
				}}
			>
				{value}
			</div>
		</div>
	);
}