import type { Agent } from "@jaeswift/jae-agent-core";
import { html, LitElement } from "lit";
import { customElement, state } from "lit/decorators.js";

export interface UsageSnapshot {
	inputTokens: number;
	outputTokens: number;
	totalTokens: number;
	estimatedCost: number;
	model: string;
	requestCount: number;
}

// Very rough cost estimates per 1M tokens for common models
const MODEL_COSTS: Record<string, { input: number; output: number }> = {
	default: { input: 3.0, output: 15.0 },
};

function estimateCost(model: string, input: number, output: number): number {
	const costs = MODEL_COSTS[model] || MODEL_COSTS.default;
	return (input / 1_000_000) * costs.input + (output / 1_000_000) * costs.output;
}

@customElement("cost-tracker")
export class CostTracker extends LitElement {
	@state() private inputTokens = 0;
	@state() private outputTokens = 0;
	@state() private requestCount = 0;
	@state() private modelId = "";
	@state() private expanded = false;

	private unsubscribe?: () => void;

	protected override createRenderRoot() {
		return this;
	}

	bindAgent(agent: Agent) {
		if (this.unsubscribe) this.unsubscribe();
		this.inputTokens = 0;
		this.outputTokens = 0;
		this.requestCount = 0;
		this.modelId = agent.state.model?.id || "";
		this.unsubscribe = agent.subscribe((event) => {
			if (event.type === "message_end" && event.message.role === "assistant") {
				const msg = event.message as any;
				if (msg.usage) {
					this.inputTokens += msg.usage.inputTokens || msg.usage.input || 0;
					this.outputTokens += msg.usage.outputTokens || msg.usage.output || 0;
					this.requestCount += 1;
				}
			}
		});
	}

	get totalTokens() {
		return this.inputTokens + this.outputTokens;
	}
	get estimatedCost() {
		return estimateCost(this.modelId, this.inputTokens, this.outputTokens);
	}

	reset() {
		this.inputTokens = 0;
		this.outputTokens = 0;
		this.requestCount = 0;
	}

	override render() {
		const cost = this.estimatedCost;
		return html`
      <button
        class="flex items-center gap-1.5 px-2 py-1 text-xs rounded hover:bg-secondary transition-colors text-muted-foreground"
        @click=${() => {
				this.expanded = !this.expanded;
				this.requestUpdate();
			}}
        title="Token usage & cost"
      >
        <span class="font-mono">${this.totalTokens > 0 ? this.totalTokens.toLocaleString() : "0"} tok</span>
        <span class="text-muted-foreground/50">|</span>
        <span class="font-mono">$${cost.toFixed(4)}</span>
      </button>
      ${
			this.expanded
				? html`
        <div class="absolute top-12 right-2 z-50 bg-popover border border-border rounded-lg shadow-xl p-4 min-w-52">
          <div class="font-semibold text-sm mb-3 flex items-center justify-between">
            <span>Token Usage</span>
            <button class="text-muted-foreground hover:text-foreground text-xs" @click=${() => this.reset()}>Reset</button>
          </div>
          <div class="flex flex-col gap-2 text-xs">
            <div class="flex justify-between"><span class="text-muted-foreground">Input tokens</span><span class="font-mono">${this.inputTokens.toLocaleString()}</span></div>
            <div class="flex justify-between"><span class="text-muted-foreground">Output tokens</span><span class="font-mono">${this.outputTokens.toLocaleString()}</span></div>
            <div class="flex justify-between font-medium border-t border-border pt-2"><span>Total tokens</span><span class="font-mono">${this.totalTokens.toLocaleString()}</span></div>
            <div class="flex justify-between text-yellow-600 dark:text-yellow-400"><span>Est. cost</span><span class="font-mono">$${cost.toFixed(6)}</span></div>
            <div class="flex justify-between"><span class="text-muted-foreground">Requests</span><span class="font-mono">${this.requestCount}</span></div>
          </div>
        </div>
      `
				: ""
		}
    `;
	}
}

export function createCostTracker(): CostTracker {
	return document.createElement("cost-tracker") as CostTracker;
}
