import React, { useEffect, useState } from 'react';
import { useApplyOp } from '../api/hooks';
import { useSessionStore } from '../store/session';
import { useUIStore } from '../store/ui';
import { useNotebookStore } from '../store/notebook';
import type { OperationInfo } from '../api/types';
import { listOperations } from '../api/client';

interface OperationCategory {
  key: string;
  label: string;
  icon: React.ReactNode;
  ops: string[]; // operation keys that belong to this category
}

const OPERATION_CATEGORIES: OperationCategory[] = [
  {
    key: 'filter_clean',
    label: 'Filter & Clean',
    icon: (
      <svg className="w-3.5 h-3.5" fill="none" viewBox="0 0 24 24" stroke="currentColor">
        <path strokeLinecap="round" strokeLinejoin="round" strokeWidth={2} d="M3 4a1 1 0 011-1h16a1 1 0 011 1v2.586a1 1 0 01-.293.707l-6.414 6.414a1 1 0 00-.293.707V17l-4 4v-6.586a1 1 0 00-.293-.707L3.293 7.293A1 1 0 013 6.586V4z" />
      </svg>
    ),
    ops: ['filter', 'fill_nulls', 'drop_columns', 'replace_values', 'distinct', 'remove_nulls', 'sample'],
  },
  {
    key: 'transform',
    label: 'Transform',
    icon: (
      <svg className="w-3.5 h-3.5" fill="none" viewBox="0 0 24 24" stroke="currentColor">
        <path strokeLinecap="round" strokeLinejoin="round" strokeWidth={2} d="M4 4v5h.582m15.356 2A8.001 8.001 0 004.582 9m0 0H9m11 11v-5h-.581m0 0a8.003 8.003 0 01-15.357-2m15.357 2H15" />
      </svg>
    ),
    ops: ['sort', 'select', 'rename', 'cast_type', 'clip', 'add_row'],
  },
  {
    key: 'aggregate',
    label: 'Aggregate',
    icon: (
      <svg className="w-3.5 h-3.5" fill="none" viewBox="0 0 24 24" stroke="currentColor">
        <path strokeLinecap="round" strokeLinejoin="round" strokeWidth={2} d="M9 19v-6a2 2 0 00-2-2H5a2 2 0 00-2 2v6m6 0V9a2 2 0 012-2h2a2 2 0 012 2v10m6 0v-4a2 2 0 00-2-2h-2a2 2 0 00-2 2v4" />
      </svg>
    ),
    ops: ['groupby', 'pivot_table', 'describe', 'value_counts'],
  },
  {
    key: 'compute',
    label: 'Compute',
    icon: (
      <svg className="w-3.5 h-3.5" fill="none" viewBox="0 0 24 24" stroke="currentColor">
        <path strokeLinecap="round" strokeLinejoin="round" strokeWidth={2} d="M9 7h6m0 10v-3m-3 3h.01M9 17h.01M9 14h.01M12 14h.01M15 11h.01M12 11h.01M9 11h.01M7 21h10a2 2 0 002-2V5a2 2 0 00-2-2H7a2 2 0 00-2 2v14a2 2 0 002 2z" />
      </svg>
    ),
    ops: ['compute_column', 'lag_lead', 'rolling_agg', 'correlation'],
  },
];

export function ManipulationPanel() {
  const { columns } = useSessionStore();
  const { addCell } = useNotebookStore();
  const { apply, loading, error } = useApplyOp();
  const { addError } = useUIStore();
  const [operations, setOperations] = useState<OperationInfo[]>([]);
  const [selectedOp, setSelectedOp] = useState<string>('');
  const [params, setParams] = useState<Record<string, any>>({});
  const [collapsedCategories, setCollapsedCategories] = useState<Set<string>>(new Set());

  useEffect(() => {
    loadOps();
  }, []);

  const loadOps = async () => {
    const data = await listOperations();
    if (data) setOperations(data.operations);
  };

  const currentOp = operations.find((op) => op.key === selectedOp);

  const handleApply = async () => {
    if (!selectedOp) return;
    const result = await apply(selectedOp, params);
    if (result) {
      addCell('table', `${selectedOp} result`);
      // Update the schema in the session store after transformation
      if (result.columns_schema) {
        useSessionStore.getState().setSchema(result.columns_schema, result.row_count);
      }
    } else if (error) {
      addError(error);
    }
  };

  const toggleCategory = (key: string) => {
    setCollapsedCategories((prev) => {
      const next = new Set(prev);
      if (next.has(key)) {
        next.delete(key);
      } else {
        next.add(key);
      }
      return next;
    });
  };

  // Categorize operations
  const categorizedOps = OPERATION_CATEGORIES.map((cat) => ({
    ...cat,
    operations: operations.filter((op) => cat.ops.includes(op.key)),
  }));

  // Find operations that don't belong to any category
  const categorizedKeys = new Set(OPERATION_CATEGORIES.flatMap((c) => c.ops));
  const uncategorizedOps = operations.filter((op) => !categorizedKeys.has(op.key));

  const renderParamField = (name: string, schema: any) => {
    const colNames = columns.map((c) => c.name);
    const numericCols = columns.filter((c) => c.column_type === 'numeric').map((c) => c.name);
    const categoricalCols = columns.filter((c) => c.column_type === 'categorical').map((c) => c.name);
    const datetimeCols = columns.filter((c) => c.column_type === 'datetime').map((c) => c.name);
    const type = schema.type;

    if (name === 'column' || name === 'x_col' || name === 'y_col') {
      const relevantCols = numericCols.length > 0 ? numericCols : colNames;
      return (
        <select
          className="input-field"
          value={params[name] || ''}
          onChange={(e) => setParams({ ...params, [name]: e.target.value })}
        >
          <option value="">Select column...</option>
          {relevantCols.map((c) => (
            <option key={c} value={c}>{c}</option>
          ))}
        </select>
      );
    }

    if (name === 'category_col' || name === 'group_col') {
      const relevantCols = categoricalCols.length > 0 ? categoricalCols : colNames;
      return (
        <select
          className="input-field"
          value={params[name] || ''}
          onChange={(e) => setParams({ ...params, [name]: e.target.value })}
        >
          <option value="">Select column...</option>
          {relevantCols.map((c) => (
            <option key={c} value={c}>{c}</option>
          ))}
        </select>
      );
    }

    if (name === 'date_col') {
      const relevantCols = datetimeCols.length > 0 ? datetimeCols : colNames;
      return (
        <select
          className="input-field"
          value={params[name] || ''}
          onChange={(e) => setParams({ ...params, [name]: e.target.value })}
        >
          <option value="">Select column...</option>
          {relevantCols.map((c) => (
            <option key={c} value={c}>{c}</option>
          ))}
        </select>
      );
    }

    if (name === 'value_col' || name === 'lat_col' || name === 'lon_col' ||
        name === 'color_col' || name === 'label_col') {
      return (
        <select
          className="input-field"
          value={params[name] || ''}
          onChange={(e) => setParams({ ...params, [name]: e.target.value })}
        >
          <option value="">Select column...</option>
          {colNames.map((c) => (
            <option key={c} value={c}>{c}</option>
          ))}
        </select>
      );
    }

    if (type === 'boolean') {
      return (
        <label className="flex items-center gap-2 text-sm">
          <input
            type="checkbox"
            checked={params[name] || false}
            onChange={(e) => setParams({ ...params, [name]: e.target.checked })}
            className="rounded border-gray-300"
          />
          {name}
        </label>
      );
    }

    if (schema.enum) {
      return (
        <select
          className="input-field"
          value={params[name] || ''}
          onChange={(e) => setParams({ ...params, [name]: e.target.value })}
        >
          <option value="">Select...</option>
          {schema.enum.map((v: string) => (
            <option key={v} value={v}>{v}</option>
          ))}
        </select>
      );
    }

    if (type === 'array') {
      return (
        <input
          className="input-field"
          placeholder={`Comma-separated ${name}`}
          value={Array.isArray(params[name]) ? params[name].join(', ') : params[name] || ''}
          onChange={(e) => setParams({ ...params, [name]: e.target.value.split(',').map((s) => s.trim()).filter(Boolean) })}
        />
      );
    }

    if (type === 'number') {
      return (
        <input
          type="number"
          className="input-field"
          placeholder={name}
          value={params[name] || ''}
          onChange={(e) => setParams({ ...params, [name]: parseFloat(e.target.value) || 0 })}
        />
      );
    }

    return (
      <input
        className="input-field"
        placeholder={name}
        value={params[name] || ''}
        onChange={(e) => setParams({ ...params, [name]: e.target.value })}
      />
    );
  };

  return (
    <div className="p-4">
      <h4 className="text-sm font-semibold text-gray-900 mb-3">Transform Data</h4>

      {/* Operation categories */}
      <div className="space-y-1 mb-4">
        {categorizedOps.map((cat) => {
          if (cat.operations.length === 0) return null;
          const isCollapsed = collapsedCategories.has(cat.key);
          return (
            <div key={cat.key}>
              <button
                onClick={() => toggleCategory(cat.key)}
                className="w-full flex items-center gap-2 py-1.5 px-2 text-xs font-semibold text-gray-600 hover:bg-gray-50 rounded-lg transition-colors"
              >
                {cat.icon}
                <span className="flex-1 text-left">{cat.label}</span>
                <svg
                  className={`w-3 h-3 text-gray-400 transition-transform ${isCollapsed ? '' : 'rotate-180'}`}
                  fill="none" viewBox="0 0 24 24" stroke="currentColor"
                >
                  <path strokeLinecap="round" strokeLinejoin="round" strokeWidth={2} d="M19 9l-7 7-7-7" />
                </svg>
              </button>
              {!isCollapsed && (
                <div className="ml-1 space-y-0.5 mt-0.5">
                  {cat.operations.map((op) => (
                    <button
                      key={op.key}
                      onClick={() => { setSelectedOp(op.key); setParams({}); }}
                      className={`op-btn ${
                        selectedOp === op.key ? 'op-btn-selected' : 'op-btn-available'
                      }`}
                    >
                      <div className="font-medium text-xs">{op.label}</div>
                      <div className="text-[10px] text-gray-400 mt-0.5">{op.description}</div>
                    </button>
                  ))}
                </div>
              )}
            </div>
          );
        })}

        {/* Uncategorized operations */}
        {uncategorizedOps.length > 0 && (
          <div>
            <button
              onClick={() => toggleCategory('_other')}
              className="w-full flex items-center gap-2 py-1.5 px-2 text-xs font-semibold text-gray-600 hover:bg-gray-50 rounded-lg transition-colors"
            >
              <svg className="w-3.5 h-3.5" fill="none" viewBox="0 0 24 24" stroke="currentColor">
                <path strokeLinecap="round" strokeLinejoin="round" strokeWidth={2} d="M5 12h.01M12 12h.01M19 12h.01M6 12a1 1 0 11-2 0 1 1 0 012 0zm7 0a1 1 0 11-2 0 1 1 0 012 0zm7 0a1 1 0 11-2 0 1 1 0 012 0z" />
              </svg>
              <span className="flex-1 text-left">Other</span>
              <svg
                className={`w-3 h-3 text-gray-400 transition-transform ${collapsedCategories.has('_other') ? '' : 'rotate-180'}`}
                fill="none" viewBox="0 0 24 24" stroke="currentColor"
              >
                <path strokeLinecap="round" strokeLinejoin="round" strokeWidth={2} d="M19 9l-7 7-7-7" />
              </svg>
            </button>
            {!collapsedCategories.has('_other') && (
              <div className="ml-1 space-y-0.5 mt-0.5">
                {uncategorizedOps.map((op) => (
                  <button
                    key={op.key}
                    onClick={() => { setSelectedOp(op.key); setParams({}); }}
                    className={`op-btn ${
                      selectedOp === op.key ? 'op-btn-selected' : 'op-btn-available'
                    }`}
                  >
                    <div className="font-medium text-xs">{op.label}</div>
                    <div className="text-[10px] text-gray-400 mt-0.5">{op.description}</div>
                  </button>
                ))}
              </div>
            )}
          </div>
        )}
      </div>

      {/* Selected operation parameters */}
      {currentOp && (
        <div className="border-t border-gray-100 pt-3 mt-3">
          <div className="flex items-center justify-between mb-2">
            <h5 className="text-xs font-semibold text-gray-700">{currentOp.label}</h5>
            <button
              onClick={() => { setSelectedOp(''); setParams({}); }}
              className="text-[10px] text-gray-400 hover:text-gray-600"
            >
              Clear
            </button>
          </div>
          <p className="text-[10px] text-gray-500 mb-3">{currentOp.description}</p>
          <div className="space-y-3 mb-4">
            {Object.entries(currentOp.param_schema.properties || {}).map(([name, schema]: [string, any]) => {
              const isRequired = currentOp.param_schema.required?.includes(name);
              return (
                <div key={name}>
                  <label className="block text-xs font-medium text-gray-600 mb-1">
                    {name} {isRequired && <span className="text-red-500">*</span>}
                  </label>
                  {renderParamField(name, schema)}
                </div>
              );
            })}
          </div>
        </div>
      )}

      {error && <p className="text-red-500 text-xs mb-3">{error}</p>}

      <button
        onClick={handleApply}
        disabled={!selectedOp || loading}
        className="btn-primary w-full"
      >
        {loading ? 'Applying...' : 'Apply'}
      </button>
    </div>
  );
}