import { getModels } from "@jaeswift/jae-ai";
import { Badge } from "@mariozechner/mini-lit/dist/Badge.js";
import { Button } from "@mariozechner/mini-lit/dist/Button.js";
import { html, LitElement, type TemplateResult } from "lit";
import { customElement, state } from "lit/decorators.js";

const TAG_COLORS: Record<string, string> = {
	tools: "bg-blue-100 text-blue-800 dark:bg-blue-900 dark:text-blue-300",
	vision: "bg-purple-100 text-purple-800 dark:bg-purple-900 dark:text-purple-300",
	reasoning: "bg-yellow-100 text-yellow-800 dark:bg-yellow-900 dark:text-yellow-300",
	code: "bg-green-100 text-green-800 dark:bg-green-900 dark:text-green-300",
	"image-generation": "bg-pink-100 text-pink-800 dark:bg-pink-900 dark:text-pink-300",
	"video-generation": "bg-red-100 text-red-800 dark:bg-red-900 dark:text-red-300",
	tts: "bg-orange-100 text-orange-800 dark:bg-orange-900 dark:text-orange-300",
	asr: "bg-teal-100 text-teal-800 dark:bg-teal-900 dark:text-teal-300",
	embedding: "bg-gray-100 text-gray-800 dark:bg-gray-700 dark:text-gray-300",
	e2ee: "bg-emerald-100 text-emerald-800 dark:bg-emerald-900 dark:text-emerald-300",
	"web-search": "bg-cyan-100 text-cyan-800 dark:bg-cyan-900 dark:text-cyan-300",
};

const CATEGORY_LABELS: Record<string, string> = {
	text: "💬 Text & Chat",
	image: "🖼️ Image Generation",
	video: "🎬 Video Generation",
	audio: "🔊 Audio (TTS / ASR)",
	other: "🔧 Other",
};

function categoriseModel(tags: string[] = []): string {
	if (tags.includes("image-generation") || tags.includes("inpainting")) return "image";
	if (tags.includes("video-generation")) return "video";
	if (tags.includes("tts") || tags.includes("asr")) return "audio";
	if (tags.includes("embedding") || tags.includes("upscaling")) return "other";
	return "text";
}

@customElement("venice-model-browser")
export class VeniceModelBrowser extends LitElement {
	@state() private filter: string = "all";
	@state() private search: string = "";

	protected createRenderRoot() {
		return this;
	}

	private renderTag(tag: string): TemplateResult {
		const cls = TAG_COLORS[tag] ?? "bg-gray-100 text-gray-700 dark:bg-gray-700 dark:text-gray-300";
		return html`<span class="inline-flex items-center px-1.5 py-0.5 rounded text-xs font-medium ${cls}">${tag}</span>`;
	}

	render(): TemplateResult {
		let models: any[] = [];
		try {
			models = (getModels("venice" as any) as any[]) || [];
		} catch {
			models = [];
		}

		// Group by category
		const grouped: Record<string, any[]> = { text: [], image: [], video: [], audio: [], other: [] };
		for (const m of models) {
			const cat = categoriseModel(m.tags);
			if (!grouped[cat]) grouped[cat] = [];
			grouped[cat].push(m);
		}

		const filters = [
			{ id: "all", label: `All (${models.length})` },
			...Object.entries(grouped)
				.filter(([, ms]) => ms.length > 0)
				.map(([cat, ms]) => ({ id: cat, label: `${CATEGORY_LABELS[cat] ?? cat} (${ms.length})` })),
		];

		const searchLower = this.search.toLowerCase();
		const activeGroups = Object.entries(grouped).filter(([cat, ms]) => {
			if (this.filter !== "all" && cat !== this.filter) return false;
			return ms.length > 0;
		});

		return html`
      <div class="flex flex-col gap-4">
        <!-- Filter bar -->
        <div class="flex flex-wrap gap-2 items-center">
          ${filters.map(
					(f) => html`
            <button
              class="px-3 py-1 rounded-full text-xs font-medium border transition-colors
                ${
							this.filter === f.id
								? "bg-primary text-primary-foreground border-primary"
								: "border-border text-muted-foreground hover:bg-secondary"
						}"
              @click=${() => {
						this.filter = f.id;
						this.requestUpdate();
					}}
            >${f.label}</button>
          `,
				)}
          <input
            type="search"
            placeholder="Search models..."
            class="ml-auto px-3 py-1 text-sm rounded border border-border bg-background text-foreground w-48"
            .value=${this.search}
            @input=${(e: Event) => {
					this.search = (e.target as HTMLInputElement).value;
					this.requestUpdate();
				}}
          />
        </div>

        <!-- Model groups -->
        ${activeGroups.map(([cat, ms]) => {
				const filtered = searchLower ? ms.filter((m: any) => m.id.toLowerCase().includes(searchLower)) : ms;
				if (!filtered.length) return html``;
				return html`
            <div class="flex flex-col gap-2">
              <h4 class="text-xs font-semibold text-muted-foreground uppercase tracking-wider">${CATEGORY_LABELS[cat] ?? cat}</h4>
              <div class="grid grid-cols-1 gap-2">
                ${filtered.map(
							(m: any) => html`
                  <div class="flex items-start justify-between p-3 rounded-lg border border-border bg-card hover:bg-secondary/40 transition-colors">
                    <div class="flex flex-col gap-1 min-w-0">
                      <span class="text-sm font-mono font-medium text-foreground truncate">${m.id}</span>
                      ${m.contextWindow ? html`<span class="text-xs text-muted-foreground">${(m.contextWindow / 1000).toFixed(0)}k ctx</span>` : ""}
                    </div>
                    <div class="flex flex-wrap gap-1 justify-end ml-2 shrink-0 max-w-48">
                      ${(m.tags ?? []).map((t: string) => this.renderTag(t))}
                    </div>
                  </div>
                `,
						)}
              </div>
            </div>
          `;
			})}

        ${
				models.length === 0
					? html`
          <div class="text-sm text-muted-foreground text-center py-4">
            No Venice models found. Ensure jae-ai includes Venice models.
          </div>
        `
					: ""
			}
      </div>
    `;
	}
}
