src/components/SearchPage.tsx

"use client";

import { useMemo } from "react";
import PageDrawer from "@/components/PageDrawer";
import ResultCard from "@/components/ResultCard";
import { exportEvidence } from "@/lib/api";
import { QUERY_MAX_CHARS, useSearchStore } from "@/lib/store";

const PRESETS = [
  {
    label: "Methods",
    query: "Summarize the experimental methods and cite the main figures.",
  },
  {
    label: "Results",
    query: "List the main findings and provide page citations.",
  },
  {
    label: "Compare",
    query: "Compare two studies in the corpus on the same topic and cite both.",
  },
  {
    label: "Definitions",
    query: "Extract key definitions and cite where they appear.",
  },
];

const formatScore = (value: number) =>
  Number.isFinite(value) ? value.toFixed(3) : "n/a";

export default function SearchPage() {
  const {
    query,
    topK,
    hybrid,
    includeImages,
    results,
    answer,
    citations,
    status,
    error,
    savedQueries,
    selectedHit,
    streamController,
    setQuery,
    setTopK,
    toggleHybrid,
    toggleImages,
    setError,
    saveQuery,
    removeSavedQuery,
    setSelectedHit,
    clearSelectedHit,
    cancelStream,
    runSearch,
    runChatStream,
  } = useSearchStore();

  const statusLabel = useMemo(() => {
    if (status === "searching") return "Retrieving evidence";
    if (status === "chatting") return "Composing answer";
    return "Idle";
  }, [status]);

  const hasResults = results.length > 0;
  const isBusy = status !== "idle";
  const queryRemaining = QUERY_MAX_CHARS - query.length;

  const downloadBlob = (blob: Blob, filename: string) => {
    const url = URL.createObjectURL(blob);
    const link = document.createElement("a");
    link.href = url;
    link.download = filename;
    link.click();
    URL.revokeObjectURL(url);
  };

  const exportJson = async () => {
    if (!query.trim()) return;
    setError(null);
    try {
      const blob = await exportEvidence({
        query,
        top_k: topK,
        include_images: includeImages,
        hybrid,
        format: "json",
        include_snippets: true,
        answer: answer || null,
      });
      downloadBlob(blob, "layra_evidence_pack.json");
    } catch (err) {
      const message = err instanceof Error ? err.message : String(err);
      setError(message || "Export failed.");
    }
  };

  const exportMarkdown = async () => {
    if (!query.trim()) return;
    setError(null);
    try {
      const blob = await exportEvidence({
        query,
        top_k: topK,
        include_images: includeImages,
        hybrid,
        format: "markdown",
        include_snippets: true,
        answer: answer || null,
      });
      downloadBlob(blob, "layra_evidence_pack.md");
    } catch (err) {
      const message = err instanceof Error ? err.message : String(err);
      setError(message || "Export failed.");
    }
  };

  return (
    <div className="min-h-screen bg-[radial-gradient(circle_at_top,_#f5efe4,_#f8f4ec_35%,_#f2ece2_60%,_#ece5da_100%)]">
      <div className="pointer-events-none absolute inset-0 overflow-hidden">
        <div className="absolute -right-32 -top-16 h-80 w-80 rounded-full bg-amber-200/50 blur-3xl" />
        <div className="absolute left-1/2 top-1/3 h-72 w-72 -translate-x-1/2 rounded-full bg-rose-200/40 blur-3xl" />
        <div className="absolute bottom-10 left-12 h-56 w-56 rounded-full bg-lime-200/30 blur-3xl" />
      </div>

      <div className="relative mx-auto flex min-h-screen w-full max-w-7xl flex-col gap-10 px-6 pb-16 pt-12">
        <header className="flex flex-col gap-4">
          <div className="flex items-center gap-3 text-xs uppercase tracking-[0.3em] text-stone-500">
            <span className="h-px w-12 bg-stone-400/60" />
            Visual RAG Studio
          </div>
          <h1 className="font-display text-4xl leading-tight text-stone-900 md:text-5xl">
            LAYRA · ColPali Research Console
          </h1>
          <p className="max-w-2xl text-base text-stone-600 md:text-lg">
            Explore academic documents with multi-vector visual retrieval, hybrid text
            signals, and citation-grounded answers.
          </p>
        </header>

        <section className="grid gap-8 lg:grid-cols-[1.05fr_1.95fr]">
          <div className="flex flex-col gap-6">
            <form
              className="flex flex-col gap-4 rounded-2xl border border-stone-200 bg-white/80 p-5 shadow-[0_30px_80px_-60px_rgba(20,20,20,0.6)] backdrop-blur"
              onSubmit={(event) => {
                event.preventDefault();
                void runSearch();
              }}
            >
              <label className="text-xs font-semibold uppercase tracking-[0.3em] text-stone-400">
                Query
              </label>
              <textarea
                value={query}
                onChange={(event) => setQuery(event.target.value)}
                placeholder="Ask about experimental methods, cite figures, or compare two papers..."
                maxLength={QUERY_MAX_CHARS}
                className="min-h-[120px] w-full resize-none rounded-2xl border border-stone-200 bg-stone-50/70 px-4 py-3 text-base text-stone-900 outline-none ring-0 transition focus:border-stone-400"
              />
              <div className="text-[11px] uppercase tracking-[0.2em] text-stone-400">
                {queryRemaining} chars left
              </div>
              <div className="flex flex-wrap items-center gap-3">
                <button
                  type="submit"
                  disabled={isBusy || !query.trim()}
                  className="rounded-full bg-stone-900 px-5 py-2 text-sm font-semibold text-white shadow-lg shadow-stone-500/30 transition hover:translate-y-[-1px] disabled:cursor-not-allowed disabled:opacity-50"
                >
                  Search
                </button>
                <button
                  type="button"
                  onClick={() => void runChatStream()}
                  disabled={isBusy || !query.trim()}
                  className="rounded-full border border-stone-400/60 px-5 py-2 text-sm font-semibold text-stone-700 transition hover:border-stone-900 hover:text-stone-900 disabled:cursor-not-allowed disabled:opacity-50"
                >
                  Ask (RAG)
                </button>
                {status === "chatting" && streamController && (
                  <button
                    type="button"
                    onClick={cancelStream}
                    className="rounded-full border border-red-200 px-4 py-2 text-xs font-semibold uppercase tracking-[0.2em] text-red-600"
                  >
                    Stop
                  </button>
                )}
                <button
                  type="button"
                  onClick={saveQuery}
                  disabled={!query.trim()}
                  className="rounded-full border border-stone-300 px-4 py-2 text-xs font-semibold uppercase tracking-[0.2em] text-stone-500 disabled:cursor-not-allowed disabled:opacity-50"
                >
                  Save Query
                </button>
                <span className="text-xs uppercase tracking-[0.25em] text-stone-400">
                  {statusLabel}
                </span>
              </div>

              <div className="grid gap-4 rounded-2xl border border-stone-100 bg-white/70 p-4">
                <div className="flex flex-wrap items-center justify-between gap-4">
                  <div className="flex items-center gap-3">
                    <button
                      type="button"
                      onClick={toggleHybrid}
                      disabled={isBusy}
                      className={`rounded-full px-4 py-1.5 text-xs font-semibold uppercase tracking-[0.2em] transition ${
                        hybrid
                          ? "bg-amber-200 text-amber-900"
                          : "bg-stone-100 text-stone-500"
                      } disabled:cursor-not-allowed disabled:opacity-50`}
                      aria-pressed={hybrid}
                    >
                      Hybrid
                    </button>
                    <button
                      type="button"
                      onClick={toggleImages}
                      disabled={isBusy}
                      className={`rounded-full px-4 py-1.5 text-xs font-semibold uppercase tracking-[0.2em] transition ${
                        includeImages
                          ? "bg-lime-200 text-lime-900"
                          : "bg-stone-100 text-stone-500"
                      } disabled:cursor-not-allowed disabled:opacity-50`}
                      aria-pressed={includeImages}
                    >
                      Images
                    </button>
                  </div>
                  <div className="flex items-center gap-3">
                    <span className="text-xs uppercase tracking-[0.2em] text-stone-500">
                      Top K
                    </span>
                    <span className="text-sm font-semibold text-stone-800">
                      {topK}
                    </span>
                  </div>
                </div>
                <input
                  type="range"
                  min={3}
                  max={20}
                  value={topK}
                  onChange={(event) => setTopK(Number(event.target.value))}
                  disabled={isBusy}
                  className="h-2 w-full cursor-pointer appearance-none rounded-full bg-stone-200 accent-stone-900"
                />
                <div className="flex flex-wrap gap-2 text-xs">
                  {PRESETS.map((preset) => (
                    <button
                      key={preset.label}
                      type="button"
                      onClick={() => setQuery(preset.query)}
                      className="rounded-full border border-stone-200 bg-white px-3 py-1 text-stone-600 transition hover:border-stone-400"
                    >
                      {preset.label}
                    </button>
                  ))}
                </div>
              </div>
            </form>

            <div className="rounded-2xl border border-stone-200 bg-white/70 p-5 backdrop-blur">
              <div className="flex items-center justify-between">
                <h2 className="font-display text-2xl text-stone-900">Evidence Summary</h2>
                <span className="text-xs uppercase tracking-[0.25em] text-stone-400">
                  {citations.length} citations
                </span>
              </div>
              {answer ? (
                <p className="mt-4 text-sm leading-relaxed text-stone-700">
                  {answer}
                </p>
              ) : (
                <p className="mt-4 text-sm text-stone-500">
                  Ask a question to generate a citation-grounded summary.
                </p>
              )}
              {citations.length > 0 && (
                <div className="mt-4 grid gap-3">
                  {citations.map((hit) => (
                    <button
                      key={`${hit.doc_id}-${hit.page_num}`}
                      type="button"
                      onClick={() => setSelectedHit(hit)}
                      className="flex items-center justify-between rounded-xl border border-stone-200 bg-stone-50/70 px-3 py-2 text-left text-xs"
                    >
                      <div className="text-stone-700">
                        <span className="font-semibold">{hit.doc_id}</span> · p.{" "}
                        {hit.page_num}
                      </div>
                      <span className="rounded-full bg-stone-200 px-2 py-1 text-[10px] uppercase tracking-[0.2em]">
                        {formatScore(hit.score)}
                      </span>
                    </button>
                  ))}
                </div>
              )}
              <div className="mt-4 flex flex-wrap gap-2 text-xs">
                <button
                  type="button"
                  onClick={exportJson}
                  disabled={!query.trim()}
                  className="rounded-full border border-stone-200 bg-white px-3 py-1 text-stone-600 disabled:cursor-not-allowed disabled:opacity-50"
                >
                  Download JSON
                </button>
                <button
                  type="button"
                  onClick={exportMarkdown}
                  disabled={!query.trim()}
                  className="rounded-full border border-stone-200 bg-white px-3 py-1 text-stone-600 disabled:cursor-not-allowed disabled:opacity-50"
                >
                  Download Markdown
                </button>
              </div>
              {error && (
                <div className="mt-4 rounded-xl border border-red-200 bg-red-50 px-3 py-2 text-xs text-red-700">
                  {error}
                </div>
              )}
            </div>

            <div className="rounded-2xl border border-stone-200 bg-white/70 p-5 backdrop-blur">
              <div className="flex items-center justify-between">
                <h2 className="font-display text-2xl text-stone-900">Saved Queries</h2>
                <span className="text-xs uppercase tracking-[0.25em] text-stone-400">
                  {savedQueries.length} saved
                </span>
              </div>
              {savedQueries.length === 0 ? (
                <p className="mt-4 text-sm text-stone-500">
                  Save commonly used queries to return quickly.
                </p>
              ) : (
                <div className="mt-4 grid gap-2">
                  {savedQueries.map((item) => (
                    <div
                      key={item}
                      className="flex items-center justify-between rounded-xl border border-stone-200 bg-white px-3 py-2 text-sm"
                    >
                      <button
                        type="button"
                        onClick={() => setQuery(item)}
                        className="text-left text-stone-700"
                      >
                        {item}
                      </button>
                      <button
                        type="button"
                        onClick={() => removeSavedQuery(item)}
                        className="rounded-full border border-stone-200 px-2 py-1 text-[10px] uppercase tracking-[0.2em] text-stone-500"
                      >
                        Remove
                      </button>
                    </div>
                  ))}
                </div>
              )}
            </div>
          </div>

          <div className="flex flex-col gap-6">
            <div className="flex items-center justify-between">
              <h2 className="font-display text-2xl text-stone-900">Retrieved Pages</h2>
              <span className="text-xs uppercase tracking-[0.25em] text-stone-400">
                {hasResults ? `${results.length} results` : "No results"}
              </span>
            </div>
            <div className="grid gap-5 sm:grid-cols-2 lg:grid-cols-3">
              {results.map((hit, index) => (
                <ResultCard
                  key={`${hit.doc_id}-${hit.page_num}`}
                  hit={hit}
                  index={index}
                  onOpen={setSelectedHit}
                />
              ))}
            </div>
            {!hasResults && (
              <div className="rounded-2xl border border-dashed border-stone-300 bg-white/40 p-8 text-center text-sm text-stone-500">
                Run a search to populate visual evidence tiles.
              </div>
            )}
          </div>
        </section>
      </div>

      <PageDrawer hit={selectedHit} onClose={clearSelectedHit} />
    </div>
  );
}