src/app/components/analysis/MessageTraceability.tsx

import { useState } from "react";
import { ChevronDown, ChevronRight, Database, Loader2, Route, Wrench } from "lucide-react";
import type { Observability, ObservabilityDataUsed } from "@/services/agenticApi";
import { formatDateTime } from "./utils";

interface MessageTraceabilityProps {
  observability?: Observability | null;
  loading?: boolean;
  error?: string | null;
  sources?: unknown[];
}

function renderUnknown(value: unknown) {
  if (value == null) return "-";
  if (typeof value === "string") return value;
  try {
    return JSON.stringify(value, null, 2);
  } catch {
    return String(value);
  }
}

function sourceLabel(source: unknown, index: number) {
  if (source && typeof source === "object") {
    const record = source as Record<string, unknown>;
    return String(record.name ?? record.title ?? record.source ?? record.source_id ?? record.id ?? `Source ${index + 1}`);
  }
  return String(source ?? `Source ${index + 1}`);
}

function DataUsedCard({ item, index }: { item: ObservabilityDataUsed; index: number }) {
  return (
    <div className="rounded-md border border-slate-200 bg-white p-2">
      <div className="flex items-center justify-between gap-2">
        <p className="font-medium text-slate-900">{item.source?.name ?? `Data source ${index + 1}`}</p>
        {item.source?.type && <span className="rounded bg-slate-100 px-1.5 py-0.5 text-[10px] uppercase text-slate-500">{item.source.type}</span>}
      </div>

      {item.tables && item.tables.length > 0 && (
        <p className="mt-1.5 text-slate-600">
          Tables: {item.tables.map((table) => table.name).join(", ")}
        </p>
      )}

      {item.columns_read && item.columns_read.length > 0 && (
        <div className="mt-1.5 flex flex-wrap gap-1">
          {item.columns_read.map((column) => (
            <span key={column.id} className="rounded border border-slate-200 bg-slate-50 px-1.5 py-0.5 text-[11px] text-slate-600">
              {column.table ? `${column.table}.${column.name}` : column.name}
            </span>
          ))}
        </div>
      )}

      {item.group_by && item.group_by.length > 0 && (
        <p className="mt-1.5 text-slate-600">Grouped by: {item.group_by.join(", ")}</p>
      )}

      {typeof item.rows_returned === "number" && (
        <p className="mt-1.5 text-slate-600">Rows returned: {item.rows_returned}</p>
      )}

      {item.query && (
        <pre className="mt-2 max-h-40 overflow-auto rounded bg-slate-950 p-2 text-[11px] leading-4 text-slate-100">{item.query}</pre>
      )}
    </div>
  );
}

export function MessageTraceability({ observability, loading, error, sources = [] }: MessageTraceabilityProps) {
  const [open, setOpen] = useState(false);
  const [detailsOpen, setDetailsOpen] = useState(false);
  const traceSources = observability?.sources?.length ? observability.sources : sources;
  const hasTrace = Boolean(observability || loading || error || traceSources.length > 0);

  if (!hasTrace) return null;

  const toolCalls = observability?.tool_calls ?? [];
  const dataUsed = observability?.data_used ?? [];

  return (
    <div className="mt-3 border-t border-slate-200 pt-3">
      <button
        type="button"
        title="Show traceability for this AI answer"
        aria-expanded={open}
        onClick={() => setOpen((current) => !current)}
        className="inline-flex items-center gap-2 rounded-md px-2 py-1 text-xs font-medium text-slate-600 hover:bg-slate-100 hover:text-slate-950"
      >
        {open ? <ChevronDown className="h-3.5 w-3.5" /> : <ChevronRight className="h-3.5 w-3.5" />}
        Traceability
        {loading && <Loader2 className="h-3.5 w-3.5 animate-spin" />}
      </button>

      {open && (
        <div className="mt-3 space-y-3 rounded-md border border-slate-200 bg-slate-50 p-3 text-xs text-slate-700">
          {loading && <p className="text-slate-500">Traceability is being prepared.</p>}
          {error && <p className="text-red-600">{error}</p>}

          {observability?.generated_at && (
            <p className="text-[11px] uppercase tracking-wide text-slate-400">Generated {formatDateTime(observability.generated_at)}</p>
          )}

          {observability?.intent && (
            <div>
              <p className="font-semibold text-slate-900">Intent</p>
              <p className="mt-1 leading-5">{observability.intent}</p>
            </div>
          )}

          {(observability?.planning || observability?.thinking) && (
            <div className="space-y-2">
              <div className="flex items-center gap-2 font-semibold text-slate-900">
                <Route className="h-3.5 w-3.5" />
                Thought and planning
              </div>
              {observability.thinking && <p className="leading-5">{observability.thinking}</p>}
              {observability.planning?.goal_restated && <p className="leading-5">{observability.planning.goal_restated}</p>}
              {observability.planning?.assumptions && observability.planning.assumptions.length > 0 && (
                <ul className="list-disc space-y-1 pl-4">
                  {observability.planning.assumptions.map((item) => <li key={item}>{item}</li>)}
                </ul>
              )}
              {observability.planning?.steps && observability.planning.steps.length > 0 && (
                <div className="space-y-1">
                  {observability.planning.steps.map((step, index) => (
                    <div key={`${step.stage ?? "step"}-${index}`} className="rounded border border-slate-200 bg-white px-2 py-1.5">
                      <p className="font-medium text-slate-800">{step.stage ?? `Step ${step.step ?? index + 1}`}</p>
                      {step.objective && <p className="mt-0.5 text-slate-600">{step.objective}</p>}
                    </div>
                  ))}
                </div>
              )}
            </div>
          )}

          <div>
            <div className="flex items-center gap-2 font-semibold text-slate-900">
              <Database className="h-3.5 w-3.5" />
              Data used
            </div>
            {dataUsed.length === 0 ? (
              <p className="mt-1 text-slate-500">No data usage recorded for this answer.</p>
            ) : (
              <div className="mt-2 space-y-2">
                {dataUsed.map((item, index) => (
                  <DataUsedCard key={`${item.source?.id ?? "source"}-${index}`} item={item} index={index} />
                ))}
              </div>
            )}
          </div>

          <div>
            <button
              type="button"
              title="Show tool call details"
              aria-expanded={detailsOpen}
              onClick={() => setDetailsOpen((current) => !current)}
              className="inline-flex items-center gap-2 rounded-md px-1.5 py-1 text-xs font-semibold text-slate-900 hover:bg-slate-100"
            >
              {detailsOpen ? <ChevronDown className="h-3.5 w-3.5" /> : <ChevronRight className="h-3.5 w-3.5" />}
              <Wrench className="h-3.5 w-3.5" />
              Details
            </button>

            {detailsOpen && (
              toolCalls.length === 0 ? (
                <p className="mt-1 pl-1.5 text-slate-500">No tool call recorded for this answer.</p>
              ) : (
                <div className="mt-2 space-y-2 pl-1.5">
                  {toolCalls.map((tool, index) => (
                    <div key={`${tool.name}-${index}`} className="rounded-md border border-slate-200 bg-white p-2">
                      <div className="flex items-center justify-between gap-2">
                        <p className="font-medium text-slate-900">{tool.name}</p>
                        {tool.status && <span className="rounded bg-slate-100 px-1.5 py-0.5 text-[10px] uppercase text-slate-500">{tool.status}</span>}
                      </div>
                      {tool.summary && <p className="mt-1 text-slate-600">{tool.summary}</p>}
                      {tool.input !== undefined && (
                        <pre className="mt-2 max-h-40 overflow-auto rounded bg-slate-950 p-2 text-[11px] leading-4 text-slate-100">{renderUnknown(tool.input)}</pre>
                      )}
                      {tool.output !== undefined && (
                        <pre className="mt-2 max-h-40 overflow-auto rounded bg-white p-2 text-[11px] leading-4 text-slate-600 ring-1 ring-slate-200">{renderUnknown(tool.output)}</pre>
                      )}
                      {tool.error && <p className="mt-2 text-red-600">{tool.error}</p>}
                    </div>
                  ))}
                </div>
              )
            )}
          </div>

          {traceSources.length > 0 && (
            <div>
              <p className="font-semibold text-slate-900">Citations and sources</p>
              <div className="mt-2 flex flex-wrap gap-1.5">
                {traceSources.map((source, index) => (
                  <span key={`${sourceLabel(source, index)}-${index}`} className="rounded-md border border-slate-200 bg-white px-2 py-1 text-[11px] text-slate-600">
                    {sourceLabel(source, index)}
                  </span>
                ))}
              </div>
            </div>
          )}
        </div>
      )}
    </div>
  );
}