import Link from "next/link"
import { unstable_cache } from "next/cache"
import { Search as SearchIcon } from "lucide-react"
import { prisma } from "@/lib/prisma"
import { safeCostTokens } from "@/lib/observability-utils"
import type { Prisma } from "@prisma/client"
import { HelpButton } from "@/components/observability/help-button"
import { StatCard } from "@/components/observability/stat-card"
import { ColumnHeader } from "@/components/observability/column-header"
import { InfoIcon } from "@/components/observability/info-icon"

// Revalidate every 30 seconds
export const revalidate = 30

interface SearchParams {
  trace_id?: string
  page?: string
  organization?: string
}

// Cached function to get unique organizations list
// Revalidates every 5 minutes - organizations change infrequently
const getUniqueOrganizations = unstable_cache(
  async () => {
    const uniqueOrganizations = await prisma.optimization_features.findMany({
      select: { organization: true },
      distinct: ["organization"],
      where: { organization: { not: null } },
    })
    return uniqueOrganizations.map(f => f.organization).filter(Boolean).sort() as string[]
  },
  ["unique-organizations"],
  { revalidate: 300 }, // 5 minutes
)

// Cached function to get trace IDs for an organization
// Shared between getTotalTracesCount and main query to avoid duplicate DB calls
const getOrgTraceIds = unstable_cache(
  async (organization: string) => {
    const orgFeatures = await prisma.optimization_features.findMany({
      where: { organization },
      select: { trace_id: true },
      distinct: ["trace_id"],
    })
    return orgFeatures.map(f => f.trace_id).filter(Boolean) as string[]
  },
  ["org-trace-ids"],
  { revalidate: 30 },
)

// Optimized function to count distinct trace_ids using groupBy
const getTotalTracesCount = unstable_cache(
  async (traceIdFilter: string | undefined, organizationFilter: string | undefined) => {
    // Get trace IDs filtered by organization if specified
    let traceIdPrefixes: string[] = []
    if (organizationFilter) {
      traceIdPrefixes = await getOrgTraceIds(organizationFilter)
      if (traceIdPrefixes.length === 0) return 0
    }

    // Build where clause
    const where: Prisma.llm_callsWhereInput = {}

    if (traceIdFilter) {
      where.trace_id = { contains: traceIdFilter }
    } else if (traceIdPrefixes.length > 0) {
      where.trace_id = { in: traceIdPrefixes }
    }

    // Use groupBy for efficient distinct count
    const result = await prisma.llm_calls.groupBy({
      by: ["trace_id"],
      where,
    })

    return result.filter(r => r.trace_id !== null).length
  },
  ["traces-count"],
  { revalidate: 30 },
)

export default async function TracesPage({ searchParams }: { searchParams: SearchParams }) {
  try {
    const page = parseInt(searchParams.page || "1")
    const pageSize = 50
    const skip = (page - 1) * pageSize

  // Get trace IDs filtered by organization if specified (uses shared cached function)
  let filteredTraceIds: string[] = []
  if (searchParams.organization) {
    filteredTraceIds = await getOrgTraceIds(searchParams.organization)

    // If organization filter is applied but no traces found, return empty result early
    if (filteredTraceIds.length === 0) {
      const orgs = await getUniqueOrganizations()

      // Return early with empty results
      return (
        <div className="container mx-auto p-6">
          {/* Header with search form */}
          <div className="mb-8">
            <div className="flex items-center justify-between gap-4 mb-2">
              <div className="flex items-center gap-3">
                <h2 className="text-2xl font-bold tracking-tight text-gray-900 dark:text-white whitespace-nowrap">
                  All Traces
                </h2>
                <HelpButton
                  title="Understanding Traces"
                  content={
                    <div className="space-y-4 text-sm text-gray-700 dark:text-gray-300">
                      <div>
                        <h4 className="font-semibold text-gray-900 dark:text-white mb-2">What is a trace?</h4>
                        <p>A trace represents a complete optimization request from start to finish. Each trace contains all the LLM API calls made during that optimization.</p>
                      </div>
                      <div>
                        <h4 className="font-semibold text-gray-900 dark:text-white mb-2">Multi-model traces</h4>
                        <p>When using multiple models for optimization, all calls share the same base trace_id (first 33 characters). This helps track related operations together.</p>
                      </div>
                      <div>
                        <h4 className="font-semibold text-gray-900 dark:text-white mb-2">Page sections</h4>
                        <ul className="list-disc list-inside space-y-1 ml-2">
                          <li><strong>Summary Stats:</strong> Quick overview of trace metrics on this page</li>
                          <li><strong>Traces Table:</strong> Detailed list of all traces with aggregated data</li>
                          <li><strong>Filters:</strong> Search by trace ID or filter by organization</li>
                        </ul>
                      </div>
                    </div>
                  }
                />
              </div>
              <form method="get" className="flex items-center gap-2 flex-1 max-w-2xl">
                <div className="flex-1 relative">
                  <input
                    type="text"
                    name="trace_id"
                    placeholder="Search by Trace ID..."
                    defaultValue={searchParams.trace_id || ""}
                    className="w-full px-3 py-2 text-sm border border-gray-300 dark:border-gray-600 rounded bg-white dark:bg-gray-700 text-gray-900 dark:text-gray-100 placeholder-gray-500 dark:placeholder-gray-400 focus:ring-2 focus:ring-blue-500 focus:border-blue-500 transition-shadow"
                    title="Search by full or partial trace ID. Finds all related calls."
                  />
                </div>
                <div className="flex items-center gap-1.5">
                  <label htmlFor="organization-filter-empty" className="text-sm font-medium text-gray-700 dark:text-gray-300 whitespace-nowrap">
                    Organization
                  </label>
                  <InfoIcon
                    content="Filter by organization that initiated the optimization request"
                    side="top"
                  />
                </div>
                <select
                  id="organization-filter-empty"
                  name="organization"
                  defaultValue={searchParams.organization || ""}
                  className="px-3 py-2 text-sm border border-gray-300 dark:border-gray-600 rounded bg-white dark:bg-gray-700 text-gray-900 dark:text-gray-100 focus:ring-2 focus:ring-blue-500 focus:border-blue-500"
                >
                  <option value="">All Organizations</option>
                  {orgs.map(org => (
                    <option key={org} value={org}>
                      {org}
                    </option>
                  ))}
                </select>
                <button
                  type="submit"
                  className="px-4 py-2 text-sm bg-blue-600 dark:bg-blue-700 text-white rounded hover:bg-blue-700 dark:hover:bg-blue-600 whitespace-nowrap"
                >
                  Search
                </button>
                <Link
                  href="/observability/traces"
                  className="px-4 py-2 text-sm border border-gray-300 dark:border-gray-600 rounded hover:bg-gray-50 dark:hover:bg-gray-700 text-gray-900 dark:text-gray-100 whitespace-nowrap"
                >
                  Clear
                </Link>
              </form>
            </div>
            <p className="text-gray-600 dark:text-gray-400">
              View optimization request traces with aggregated metrics
            </p>
          </div>
          <div className="bg-white dark:bg-gray-800 rounded-lg shadow p-16 text-center border border-gray-200 dark:border-gray-700">
            <div className="flex flex-col items-center justify-center space-y-4">
              <SearchIcon className="h-12 w-12 text-gray-400 dark:text-gray-500" />
              <div>
                <h3 className="text-lg font-semibold text-gray-900 dark:text-white mb-1">
                  No Traces Found
                </h3>
                <p className="text-sm text-gray-500 dark:text-gray-400">
                  No traces found for organization &quot;{searchParams.organization}&quot;. Try selecting a different organization.
                </p>
              </div>
              <Link
                href="/observability/llm-calls"
                className="text-sm text-blue-600 dark:text-blue-400 hover:underline"
              >
                View All LLM Calls →
              </Link>
            </div>
          </div>
        </div>
      )
    }
  }

  // Build where clause for LLM calls with organization filter applied at database level
  // NOTE: Filtering happens at DB level BEFORE pagination, not client-side.
  // We use IN clause because there's no Prisma relation between llm_calls and optimization_features
  // (they're only related by trace_id as a string field, not a foreign key relation)
  const where: Prisma.llm_callsWhereInput = {}

  if (searchParams.trace_id) {
    where.trace_id = { contains: searchParams.trace_id }
  } else if (filteredTraceIds.length > 0) {
    // Use IN clause for exact trace ID matches - much more efficient than OR with startsWith
    // For very large organizations (>10k traces), consider chunking the array
    where.trace_id = { in: filteredTraceIds }
  }

    // STEP 1: Get distinct trace_ids with pagination using groupBy
    const [distinctTraces, totalTracesCount] = await Promise.all([
      prisma.llm_calls.groupBy({
        by: ["trace_id"],
        where,
        orderBy: { _max: { created_at: "desc" } },
        take: pageSize,
        skip,
        _max: { created_at: true },
      }),
      getTotalTracesCount(searchParams.trace_id, searchParams.organization),
    ])

    // Extract trace_ids from the paginated results
    const paginatedTraceIds = distinctTraces
      .map(t => t.trace_id)
      .filter(Boolean) as string[]

    // STEP 2: Fetch all LLM calls ONLY for the paginated trace_ids
    const llmCallsRaw = paginatedTraceIds.length > 0
      ? await prisma.llm_calls.findMany({
          where: { trace_id: { in: paginatedTraceIds } },
          orderBy: { created_at: "desc" },
          select: {
            trace_id: true,
            created_at: true,
            llm_cost: true,
            total_tokens: true,
            status: true,
            call_type: true,
          },
        })
      : []

    // Filter out null trace_ids
    const llmCalls = llmCallsRaw.filter(call => call.trace_id !== null)

  // Fetch organizations ONLY for the paginated trace_ids
  const [allOptimizationFeatures] = await Promise.all([
    paginatedTraceIds.length > 0
      ? prisma.optimization_features.findMany({
          where: { trace_id: { in: paginatedTraceIds } },
          select: {
            trace_id: true,
            organization: true,
          },
        })
      : [],
  ])

  // Create a map of trace_id to organization
  const traceIdToOrganization = new Map<string, string>()
  allOptimizationFeatures.forEach(feature => {
    if (feature.trace_id && feature.organization) {
      traceIdToOrganization.set(feature.trace_id, feature.organization)
    }
  })

  // Get unique organizations for filter dropdown (cached)
  const orgs = await getUniqueOrganizations()

  // Group by trace_id and calculate aggregates
  const traceMap = new Map<
    string,
    {
      trace_id: string
      first_seen: Date
      last_seen: Date
      call_count: number
      total_cost: number
      total_tokens: number
      failed_calls: number
      hasPartial: boolean
      status: string
      call_types: Set<string>
    }
  >()

  llmCalls.forEach(call => {
    if (!call.trace_id) return

    const callTimestamp = new Date(call.created_at).getTime()
    const existing = traceMap.get(call.trace_id)
    const { cost: callCost, tokens: callTokens } = safeCostTokens(call.llm_cost, call.total_tokens)

    if (existing) {
      existing.call_count++
      existing.total_cost += callCost
      existing.total_tokens += callTokens
      if (call.status === "failed") existing.failed_calls++
      if (call.status === "partial_success") existing.hasPartial = true
      if (call.call_type) existing.call_types.add(call.call_type)
      existing.first_seen = new Date(Math.min(existing.first_seen.getTime(), callTimestamp))
      existing.last_seen = new Date(Math.max(existing.last_seen.getTime(), callTimestamp))
    } else {
      traceMap.set(call.trace_id, {
        trace_id: call.trace_id,
        first_seen: new Date(callTimestamp),
        last_seen: new Date(callTimestamp),
        call_count: 1,
        total_cost: callCost,
        total_tokens: callTokens,
        failed_calls: call.status === "failed" ? 1 : 0,
        hasPartial: call.status === "partial_success",
        status: call.status || "unknown",
        call_types: new Set(call.call_type ? [call.call_type] : []),
      })
    }
  })

    // Convert to array and sort by last_seen desc
    // No need to paginate here - already done at database level
    const traces = Array.from(traceMap.values()).sort(
      (a, b) => b.last_seen.getTime() - a.last_seen.getTime(),
    )
    const totalPages = Math.ceil(totalTracesCount / pageSize)

    return (
    <div className="container mx-auto p-6">
      {/* Header */}
      <div className="mb-8">
        {/* Title, Search Bar, and Help Button */}
        <div className="flex items-center justify-between gap-4 mb-2">
          <div className="flex items-center gap-3">
            <h2 className="text-2xl font-bold tracking-tight text-gray-900 dark:text-white whitespace-nowrap">
              All Traces
            </h2>
            <HelpButton
              title="Understanding Traces"
              content={
                <div className="space-y-4 text-sm text-gray-700 dark:text-gray-300">
                  <div>
                    <h4 className="font-semibold text-gray-900 dark:text-white mb-2">What is a trace?</h4>
                    <p>A trace represents a complete optimization request from start to finish. Each trace contains all the LLM API calls made during that optimization.</p>
                  </div>
                  <div>
                    <h4 className="font-semibold text-gray-900 dark:text-white mb-2">Multi-model traces</h4>
                    <p>When using multiple models for optimization, all calls share the same base trace_id (first 33 characters). This helps track related operations together.</p>
                  </div>
                  <div>
                    <h4 className="font-semibold text-gray-900 dark:text-white mb-2">Page sections</h4>
                    <ul className="list-disc list-inside space-y-1 ml-2">
                      <li><strong>Summary Stats:</strong> Quick overview of trace metrics on this page</li>
                      <li><strong>Traces Table:</strong> Detailed list of all traces with aggregated data</li>
                      <li><strong>Filters:</strong> Search by trace ID or filter by organization</li>
                    </ul>
                  </div>
                </div>
              }
            />
          </div>
          {/* Compact Search Bar */}
          <form method="get" className="flex items-center gap-2 flex-1 max-w-2xl">
            <div className="flex-1 relative">
              <input
                type="text"
                name="trace_id"
                placeholder="Search by Trace ID..."
                defaultValue={searchParams.trace_id || ""}
                className="w-full px-3 py-2 text-sm border border-gray-300 dark:border-gray-600 rounded bg-white dark:bg-gray-700 text-gray-900 dark:text-gray-100 placeholder-gray-500 dark:placeholder-gray-400 focus:ring-2 focus:ring-blue-500 focus:border-blue-500 transition-shadow"
                title="Search by full or partial trace ID. Finds all related calls."
              />
            </div>
            <div className="flex items-center gap-1.5">
              <label htmlFor="organization-filter" className="text-sm font-medium text-gray-700 dark:text-gray-300 whitespace-nowrap">
                Organization
              </label>
              <InfoIcon
                content="Filter by organization that initiated the optimization request"
                side="top"
              />
            </div>
            <select
              id="organization-filter"
              name="organization"
              defaultValue={searchParams.organization || ""}
              className="px-3 py-2 text-sm border border-gray-300 dark:border-gray-600 rounded bg-white dark:bg-gray-700 text-gray-900 dark:text-gray-100 focus:ring-2 focus:ring-blue-500 focus:border-blue-500"
            >
              <option value="">All Organizations</option>
              {orgs.map(org => (
                <option key={org} value={org}>
                  {org}
                </option>
              ))}
            </select>
            <button
              type="submit"
              className="px-4 py-2 text-sm bg-blue-600 dark:bg-blue-700 text-white rounded hover:bg-blue-700 dark:hover:bg-blue-600 whitespace-nowrap"
            >
              Search
            </button>
            {(searchParams.trace_id || searchParams.organization) && (
              <Link
                href="/observability/traces"
                className="px-4 py-2 text-sm border border-gray-300 dark:border-gray-600 rounded hover:bg-gray-50 dark:hover:bg-gray-700 text-gray-900 dark:text-gray-100 whitespace-nowrap"
              >
                Clear
              </Link>
            )}
          </form>
        </div>
        <p className="text-gray-600 dark:text-gray-400">
          View optimization request traces with aggregated metrics
        </p>
      </div>

      {/* Summary Stats */}
      <div className="grid grid-cols-1 md:grid-cols-4 gap-5 mb-8">
        <StatCard
          label="Total Traces"
          value={totalTracesCount}
          helpText="Total number of unique traces across all pages matching your filters"
          icon="Database"
        />
        <StatCard
          label="LLM Calls (Page)"
          value={llmCalls.length}
          helpText="Number of individual LLM API calls on this page. Each trace may have multiple calls."
          icon="Zap"
          variant="default"
          className="border-l-4 border-blue-500"
        />
        <StatCard
          label="Total Cost (Page)"
          value={`$${traces.reduce((sum, t) => sum + t.total_cost, 0).toFixed(4)}`}
          helpText="Total cost in USD for all LLM calls shown on this page. Based on model pricing and token usage."
          icon="DollarSign"
        />
        <StatCard
          label="Failed Traces (Page)"
          value={traces.filter(t => t.failed_calls > 0).length}
          helpText="Traces containing at least one failed LLM call. Click a trace to see error details."
          icon="AlertTriangle"
          variant="error"
        />
      </div>

      {/* Traces Table */}
      <div className="bg-white dark:bg-gray-800 rounded-lg shadow border border-gray-200 dark:border-gray-700">
        <div className="overflow-x-auto">
          <table className="w-full">
            <thead className="bg-gray-50 dark:bg-gray-900">
              <tr>
                <ColumnHeader
                  label="Trace ID"
                  tooltip="Unique identifier for this optimization request. Click to view detailed timeline."
                />
                <ColumnHeader
                  label="Organization"
                  tooltip="Organization that initiated this optimization (from GitHub/CLI metadata)"
                />
                <ColumnHeader
                  label="Status"
                  tooltip="Overall status: Success (all calls succeeded), Partial (some succeeded), Failed (any call failed)"
                />
                <ColumnHeader
                  label="Calls"
                  tooltip="Number of LLM API calls in this trace. Includes retries and multi-model attempts."
                />
                <ColumnHeader
                  label="Call Types"
                  tooltip="Types of LLM operations: optimization, validation, line_profiler, etc."
                />
                <ColumnHeader
                  label="Cost"
                  tooltip="Total cost in USD for all LLM calls in this trace"
                />
                <ColumnHeader
                  label="Tokens"
                  tooltip="Combined prompt + completion tokens across all calls"
                />
                <ColumnHeader
                  label="Duration"
                  tooltip="Time between first and last call in this trace"
                />
                <ColumnHeader
                  label="Last Seen"
                  tooltip="Timestamp of the most recent call in this trace"
                />
              </tr>
            </thead>
            <tbody className="divide-y divide-gray-200 dark:divide-gray-700">
              {traces.length === 0 ? (
                <tr>
                  <td
                    colSpan={9}
                    className="px-6 py-16 text-center"
                  >
                    <div className="flex flex-col items-center justify-center space-y-4">
                      <SearchIcon className="h-12 w-12 text-gray-400 dark:text-gray-500" />
                      <div>
                        <h3 className="text-lg font-semibold text-gray-900 dark:text-white mb-1">
                          No Traces Found
                        </h3>
                        <p className="text-sm text-gray-500 dark:text-gray-400">
                          {searchParams.trace_id || searchParams.organization
                            ? "Try adjusting your filters above"
                            : "Run an optimization to see traces here"}
                        </p>
                      </div>
                      {(searchParams.trace_id || searchParams.organization) && (
                        <Link
                          href="/observability/llm-calls"
                          className="text-sm text-blue-600 dark:text-blue-400 hover:underline"
                        >
                          View All LLM Calls →
                        </Link>
                      )}
                    </div>
                  </td>
                </tr>
              ) : (
                traces.map(trace => {
                  // Ensure duration is never negative by using Math.max
                  const duration = Math.max(
                    0,
                    (trace.last_seen.getTime() - trace.first_seen.getTime()) / 1000,
                  )
                  const hasPartial = trace.hasPartial
                  const statusColor =
                    trace.failed_calls > 0
                      ? "bg-red-100 text-red-800 dark:bg-red-900 dark:text-red-300"
                      : hasPartial
                        ? "bg-yellow-100 text-yellow-800 dark:bg-yellow-900 dark:text-yellow-300"
                        : "bg-green-100 text-green-800 dark:bg-green-900 dark:text-green-300"
                  const statusText = trace.failed_calls > 0 ? "Failed" : hasPartial ? "Partial" : "Success"
                  const statusBorderColor = trace.failed_calls > 0 ? "border-red-500" : hasPartial ? "border-yellow-500" : "border-green-500"

                  return (
                    <tr
                      key={trace.trace_id}
                      className={`border-l-4 border-transparent hover:bg-gray-50 dark:hover:bg-gray-700/50 hover:shadow-md hover:${statusBorderColor} transition-all duration-200`}
                    >
                      <td className="px-6 py-5">
                        {trace.trace_id && trace.trace_id.trim() ? (
                          <Link
                            href={`/observability/trace/${trace.trace_id}`}
                            className="text-blue-600 dark:text-blue-400 hover:underline font-mono text-sm"
                          >
                            {trace.trace_id.substring(0, 8)}...
                          </Link>
                        ) : (
                          <span className="text-gray-400 dark:text-gray-500 font-mono text-sm">N/A</span>
                        )}
                      </td>
                      <td className="px-6 py-5 whitespace-nowrap text-sm text-gray-700 dark:text-gray-300">
                        {traceIdToOrganization.get(trace.trace_id) || "N/A"}
                      </td>
                      <td className="px-6 py-5 whitespace-nowrap">
                        <span className={`inline-flex px-2 py-1 text-xs font-semibold rounded-full ${statusColor}`}>
                          {statusText}
                        </span>
                      </td>
                      <td className="px-6 py-5 whitespace-nowrap text-sm text-gray-700 dark:text-gray-300">
                        {trace.call_count}
                        {trace.failed_calls > 0 && (
                          <span className="text-red-600 dark:text-red-400 ml-1">
                            ({trace.failed_calls} failed)
                          </span>
                        )}
                      </td>
                      <td className="px-6 py-5 text-sm text-gray-700 dark:text-gray-300">
                        <div className="flex flex-wrap gap-1">
                          {Array.from(trace.call_types)
                            .slice(0, 3)
                            .map(type => (
                              <span
                                key={type}
                                className="px-2 py-0.5 text-xs bg-gray-100 dark:bg-gray-700 rounded"
                              >
                                {type}
                              </span>
                            ))}
                          {trace.call_types.size > 3 && (
                            <span className="px-2 py-0.5 text-xs text-gray-500 dark:text-gray-400">
                              +{trace.call_types.size - 3}
                            </span>
                          )}
                        </div>
                      </td>
                      <td className="px-6 py-5 whitespace-nowrap text-sm text-gray-700 dark:text-gray-300">
                        ${trace.total_cost.toFixed(4)}
                      </td>
                      <td className="px-6 py-5 whitespace-nowrap text-sm text-gray-700 dark:text-gray-300">
                        {trace.total_tokens.toLocaleString()}
                      </td>
                      <td className="px-6 py-5 whitespace-nowrap text-sm text-gray-700 dark:text-gray-300">
                        {duration.toFixed(2)}s
                      </td>
                      <td className="px-6 py-5 whitespace-nowrap text-sm text-gray-600 dark:text-gray-400">
                        {trace.last_seen.toLocaleString()}
                      </td>
                    </tr>
                  )
                })
              )}
            </tbody>
          </table>
        </div>
      </div>

      {/* Pagination */}
      {totalPages > 1 && (
        <div className="mt-8 flex justify-center items-center gap-4">
          {page > 1 && (
            <Link
              href={`/observability/traces?page=${page - 1}${searchParams.trace_id ? `&trace_id=${searchParams.trace_id}` : ""}${searchParams.organization ? `&organization=${searchParams.organization}` : ""}`}
              className="px-4 py-2 border border-gray-300 dark:border-gray-600 rounded hover:bg-gray-50 dark:hover:bg-gray-700 text-gray-900 dark:text-gray-100"
            >
              Previous
            </Link>
          )}
          <span className="text-gray-600 dark:text-gray-400">
            Page {page} of {totalPages} • Showing {traces.length} traces
          </span>
          {page < totalPages && (
            <Link
              href={`/observability/traces?page=${page + 1}${searchParams.trace_id ? `&trace_id=${searchParams.trace_id}` : ""}${searchParams.organization ? `&organization=${searchParams.organization}` : ""}`}
              className="px-4 py-2 border border-gray-300 dark:border-gray-600 rounded hover:bg-gray-50 dark:hover:bg-gray-700 text-gray-900 dark:text-gray-100"
            >
              Next
            </Link>
          )}
        </div>
      )}
    </div>
    )
  } catch (error) {
    throw error
  }
}
