import styles from "./ExperimentLogs.module.css";
import {LogMessages} from "../../../components/Logging/Logging.tsx";
import {useEffect, useMemo, useState} from "react";
import {type LogFilterPredicate, type LogRecord, useLogs} from "../../../components/Logging/useLogs.ts";
import capitalize from "../../../utils/capitalize.ts";
import {useCell} from "../../../utils/cellStore.ts";
import {
  EXPERIMENT_FILTER_KEY,
  EXPERIMENT_LOGGER_NAME,
  type LoggingSettings,
  type MessageComponentProps,
} from "../../../components/Logging/Definitions.ts";
import formatDuration from "../../../utils/formatDuration.ts";
import {create} from "zustand";
import Dialog from "../../../components/Dialog.tsx";
import delayedResolve from "../../../utils/delayedResolve.ts";

/**
 * Local Zustand store for logging UI preferences.
 */
const useLoggingSettings = create<LoggingSettings>((set) => ({
  showRelativeTime: false,
  setShowRelativeTime: (showRelativeTime: boolean) => set({ showRelativeTime }),
}));

/**
 * A dedicated component for rendering chat messages.
 *
 * @param record The chat record to render.
 */
function ChatMessage({ record }: { record: LogRecord }) {
  const { showRelativeTime, setShowRelativeTime } = useLoggingSettings();

  const reverse = record.role === "user" ? styles.alternate : "";

  return <div className={`${styles.chatMessage} ${reverse} flex-col padding-md padding-h-lg shadow-md round-md`}>
    <div className={`${styles.messageHead} flex-row gap-md align-center`}>
      <span className={"bold"}>{capitalize(record.role ?? "unknown")}</span>
      <span className={"font-small"}>•</span>
      <span className={"mono clickable font-small"} onClick={() => setShowRelativeTime(!showRelativeTime)}>{
        showRelativeTime
          ? formatDuration(record.relativeCreated)
          : new Date(record.created * 1000).toLocaleTimeString()
      }</span>
    </div>
    <span>{record.message}</span>
  </div>
}

/**
 * A generic log message component showing the log level, time, and message text.
 *
 * @param record The log record to render.
 */
function DefaultMessage({ record }: { record: LogRecord }) {
  const { showRelativeTime, setShowRelativeTime } = useLoggingSettings();

  return <div>
    <div className={"flex-row gap-md align-center"}>
      <span className={"font-small"}>{record.levelname}</span>
      <span className={"font-small"}>•</span>
      <span className={"mono clickable font-small"} onClick={() => setShowRelativeTime(!showRelativeTime)}>{
        showRelativeTime
          ? formatDuration(record.relativeCreated)
          : new Date(record.created * 1000).toLocaleTimeString()
      }</span>
    </div>
    <span>{record.message}</span>
  </div>;
}

/**
 * A custom component for rendering experiment messages, which might include chat messages.
 *
 * @param recordCell The cell containing the log record to render.
 * @param onUpdate A callback to notify the parent component when the record changes.
 */
function ExperimentMessage({recordCell, onUpdate}: MessageComponentProps) {
  const record = useCell(recordCell);

  // Notify the parent component (e.g., for scroll updates) when this record changes.
  useEffect(() => {
    if (onUpdate) onUpdate();
  }, [record, onUpdate]);

  if (record.levelname == "CHAT") {
    return <ChatMessage record={record} />
  } else {
    return <DefaultMessage record={record} />
  }
}

/**
 * A download dialog listing experiment logs to download.
 *
 * @param filenames The list of available experiment logs to download.
 * @param refresh A callback to refresh the list of available experiment logs.
 */
function DownloadScreen({filenames, refresh}: {filenames: string[] | null, refresh: () => void}) {
  const list = (() => {
    if (filenames == null) return <div className={`${styles.downloadList} flex-col align-center justify-center`}>
      <p>Loading...</p>
    </div>;
    if (filenames.length === 0) return <div className={`${styles.downloadList} flex-col align-center justify-center`}>
      <p>No files available.</p>
    </div>

    return <ol className={`${styles.downloadList} margin-0 padding-h-lg scroll-y`}>
      {filenames!.map((filename) => (
        <li><a key={filename} href={`http://localhost:8000/api/logs/files/${filename}`} download={filename}>{filename}</a></li>
      ))}
    </ol>;
  })();

  return <div className={"flex-col"}>
    <p className={"margin-lg"}>Select a file to download:</p>
    {list}
    <button onClick={refresh} className={"margin-lg shadow-sm"}>Refresh</button>
  </div>;
}

/**
 * A button that opens a download dialog for experiment logs when pressed.
 */
function DownloadButton() {
  const [showModal, setShowModal] = useState(false);
  const [filenames, setFilenames] = useState<string[] | null>(null);

  async function getFiles(): Promise<string[]> {
    const response = await fetch("http://localhost:8000/api/logs/files");
    const files = await response.json();
    files.sort();
    return files;
  }

  useEffect(() => {
    getFiles().then(setFilenames);
  }, [showModal]);

  return <>
    <button className={"shadow-sm"} onClick={() => setShowModal((curr) => !curr)}>Download...</button>
    <Dialog open={showModal} close={() => setShowModal(false)} classname={"padding-0 round-lg"}>
      <DownloadScreen filenames={filenames} refresh={async () => {
        setFilenames(null);
        const files = await delayedResolve(getFiles(), 250);
        setFilenames(files);
      }} />
    </Dialog>
  </>;
}

/**
 * A component for rendering experiment logs. This component uses the `useLogs` hook with a filter to show only
 * experiment logs.
 */
export default function ExperimentLogs() {
  // Show only experiment logs in this logger
  const filters = useMemo(() => new Map<string, LogFilterPredicate>([
    [
      EXPERIMENT_FILTER_KEY,
      {
        predicate: (r) => r.name == EXPERIMENT_LOGGER_NAME,
        priority: 999,
        value: null,
      } as LogFilterPredicate,
    ],
  ]), []);

  const { filteredLogs } = useLogs(filters);

  return <aside className={`${styles.logs} flex-col relative`}>
    <div className={`${styles.head} padding-lg`}>
      <div className={"flex-row align-center justify-between"}>
        <h3>Logs</h3>
        <div className={"flex-row gap-md align-center"}>
          <div className={`flex-row align-center gap-md relative padding-md shadow-sm round-md`}>
            <div className={styles.live}></div>
            <span>Live</span>
          </div>
          <DownloadButton />
        </div>
      </div>
    </div>
    <LogMessages recordCells={filteredLogs} MessageComponent={ExperimentMessage} />
  </aside>;
}