import { useEffect, useMemo, useState } from "react";
import { vscode } from "../utils/vscode";
import { useStore } from "../store/root";
import { canRetryTaskInQueue } from "@codeflash/shared";
import styles from "../styles/currentFileFunctions.module.css";

type FunctionItem = {
  fullName: string;
  methodName?: string;
};

const CurrentFileFunctions = () => {
  const functions = useStore((state) => state.functionsInCurrentFile);
  const queueTasks = useStore((state) => state.queueTasks);
  const currentFilePath = useStore((state) => state.currentFilePath);
  const [expanded, setExpanded] = useState(
    queueTasks.length == 0 ? true : false,
  );

  const [disabledFunctions, setDisabledFunctions] = useState<string[]>([]);

  useEffect(() => {
    const disabledFunctions = queueTasks
      .filter(
        (task) =>
          task.filepath === currentFilePath && !canRetryTaskInQueue(task),
      )
      .map((task) => task.functionName);
    setDisabledFunctions(disabledFunctions);
  }, [queueTasks, currentFilePath]);

  const { sortedClassNames, classMethods, standaloneFunctions } =
    useMemo(() => {
      const classMethods: Record<string, FunctionItem[]> = {};
      const standaloneFunctions: string[] = [];

      functions.forEach((funcName) => {
        const dotIndex = funcName.indexOf(".");
        if (dotIndex > 0 && dotIndex < funcName.length - 1) {
          const className = funcName.substring(0, dotIndex);
          const methodName = funcName.substring(dotIndex + 1);
          if (!classMethods[className]) {
            classMethods[className] = [];
          }
          classMethods[className].push({ fullName: funcName, methodName });
        } else {
          standaloneFunctions.push(funcName);
        }
      });

      const sortedClassNames = Object.keys(classMethods).sort((a, b) =>
        a.localeCompare(b, undefined, { sensitivity: "base" }),
      );

      Object.values(classMethods).forEach((methods) =>
        methods.sort((a, b) =>
          (a.methodName ?? "").localeCompare(b.methodName ?? "", undefined, {
            sensitivity: "base",
          }),
        ),
      );

      standaloneFunctions.sort((a, b) =>
        a.localeCompare(b, undefined, { sensitivity: "base" }),
      );

      return { sortedClassNames, classMethods, standaloneFunctions };
    }, [functions]);

  return (
    <div className={styles.functionsWrapper}>
      {/* Header */}
      <div
        className={styles.sectionHeader}
        onClick={(e) => {
          e.stopPropagation();
          if (functions.length === 0) {
            return;
          }
          setExpanded((prev) => !prev);
        }}
      >
        <div className={styles.sectionTitle}>
          <span className={styles.sectionTitleIcon}>⚡</span>
          <span>Optimizable Functions</span>
          {functions.length > 0 && (
            <span
              id="function-count"
              className={styles.badge}
              title="Number of optimizable functions found"
            >
              {functions.length}
            </span>
          )}
          <button
            id="refresh-button-secondary"
            title="Refresh Analysis"
            className="action-button refresh-inline"
            onClick={(e) => {
              e.stopPropagation();
              vscode.postMessage({
                type: "requestAnalysis",
              });
            }}
          >
            <span className="codicon codicon-refresh"></span>
          </button>
        </div>
      </div>

      {/* Function List */}
      <ul
        className={
          styles.functionList +
          " " +
          (expanded ? styles.expanded : styles.collapsed)
        }
      >
        {sortedClassNames.map((className) => (
          <ClassSection
            key={className}
            className={className}
            methods={classMethods[className] ?? []}
            disabledFunctions={disabledFunctions}
          />
        ))}

        {sortedClassNames.length > 0 && standaloneFunctions.length > 0 && (
          <li className="list-separator"></li>
        )}

        {standaloneFunctions.map((fn) => (
          <FunctionItemRow
            key={fn}
            fullName={fn}
            isDisabled={disabledFunctions.includes(fn)}
          />
        ))}
      </ul>
    </div>
  );
};

const ClassSection = ({
  className,
  methods,
  disabledFunctions,
}: {
  className: string;
  methods: FunctionItem[];
  disabledFunctions: string[];
}) => {
  const [expanded, setExpanded] = useState(false);

  return (
    <li
      className={styles.classSection + " " + (expanded ? styles.expanded : "")}
    >
      <div
        className={`${styles.classHeader}`}
        title={`Class: ${className}`}
        onClick={(e) => {
          if ((e.target as HTMLElement).closest("button")) {
            return;
          }
          setExpanded((prev) => !prev);
        }}
      >
        <span
          className={`class-toggle codicon codicon-chevron-right ${styles.classToggle}`}
        ></span>
        <span
          className={`codicon codicon-symbol-class ${styles.classIcon}`}
        ></span>
        <span
          className={styles.className}
          data-full-name={className}
          title={className}
        >
          {className}
        </span>
        <span
          className={`${styles.methodCount} badge`}
          title={`${methods.length} method${methods.length !== 1 ? "s" : ""}`}
        >
          {methods.length}
        </span>
      </div>
      {expanded && (
        <ul className={styles.methodsList}>
          {methods.map((m) => (
            <FunctionItemRow
              key={m.fullName}
              fullName={m.fullName}
              methodName={m.methodName ?? ""}
              isMethod
              isDisabled={disabledFunctions.includes(m.fullName)}
            />
          ))}
        </ul>
      )}
    </li>
  );
};

const FunctionItemRow = ({
  fullName,
  methodName,
  isMethod = false,
  isDisabled = false,
}: {
  fullName: string;
  methodName?: string;
  isMethod?: boolean;
  isDisabled: boolean;
}) => {
  const handleClick = () => {
    vscode.postMessage({
      type: "navigateToFunction",
      payload: { functionName: fullName },
    });
  };

  const handleOptimize = (e: React.MouseEvent) => {
    e.stopPropagation();
    vscode.postMessage({
      type: "optimizeFunction",
      payload: { functionName: fullName },
    });
  };

  return (
    <li
      className={[
        styles.functionItem,
        isMethod ? styles.methodItem : styles.standaloneFunctionItem,
        isDisabled ? styles.disabled : "",
      ].join(" ")}
      data-function={fullName}
    >
      <div
        className={styles.functionItemContent}
        title={`Navigate to: ${fullName}`}
        onClick={handleClick}
      >
        <span
          className={`codicon ${
            isMethod ? "codicon-symbol-method" : "codicon-symbol-function"
          } ${styles.functionIcon}`}
        ></span>
        <span
          className={styles.functionName}
          data-full-name={fullName}
          title={fullName}
        >
          {methodName ?? fullName}
        </span>
      </div>
      <div className="function-buttons">
        <button
          className={styles.optimizeButton}
          title={`Optimize: ${fullName}`}
          onClick={handleOptimize}
          disabled={isDisabled}
        >
          Optimize
        </button>
      </div>
    </li>
  );
};

export default CurrentFileFunctions;
