import { describe, it, expect, beforeEach, vi } from 'vitest';
import { TFJSModelManager, TensorUtils, createModelManager } from './tfjs';
import * as tf from '@tensorflow/tfjs';

// Mock TensorFlow.js
vi.mock('@tensorflow/tfjs', () => ({
  setBackend: vi.fn().mockResolvedValue(undefined),
  ready: vi.fn().mockResolvedValue(undefined),
  getBackend: vi.fn().mockReturnValue('webgl'),
  loadLayersModel: vi.fn(),
  loadGraphModel: vi.fn(),
  tidy: vi.fn((fn) => fn()),
  tensor: vi.fn((data) => ({
    shape: [data.length],
    dispose: vi.fn(),
    data: vi.fn().mockResolvedValue(new Float32Array(data)),
    arraySync: vi.fn().mockReturnValue(data),
    argMax: vi.fn().mockReturnValue({ data: vi.fn().mockResolvedValue([0]), dispose: vi.fn() }),
    abs: vi.fn().mockReturnThis(),
    max: vi.fn().mockReturnValue(1),
    div: vi.fn().mockReturnThis(),
    mean: vi.fn().mockReturnValue(0),
    sub: vi.fn().mockReturnThis(),
    add: vi.fn().mockReturnThis(),
    sqrt: vi.fn().mockReturnThis(),
    reshape: vi.fn().mockReturnThis(),
  })),
  stack: vi.fn((tensors) => ({
    dispose: vi.fn(),
    arraySync: vi.fn().mockReturnValue([[1, 2, 3]]),
  })),
  moments: vi.fn(() => ({
    variance: {
      sqrt: vi.fn().mockReturnValue({
        add: vi.fn().mockReturnValue(1),
      }),
    },
  })),
  softmax: vi.fn((t) => t),
  topk: vi.fn((tensor, k) => ({
    indices: { data: vi.fn().mockResolvedValue([0, 1, 2]), dispose: vi.fn() },
    values: { data: vi.fn().mockResolvedValue([0.9, 0.8, 0.7]), dispose: vi.fn() },
  })),
  memory: vi.fn(() => ({
    numTensors: 0,
    numDataBuffers: 0,
    numBytes: 0,
  })),
}));

describe('TFJSModelManager', () => {
  let manager: TFJSModelManager;

  beforeEach(() => {
    manager = new TFJSModelManager();
    vi.clearAllMocks();
  });

  describe('initialize', () => {
    it('should initialize with webgl backend', async () => {
      await manager.initialize('webgl');
      
      expect(tf.setBackend).toHaveBeenCalledWith('webgl');
      expect(tf.ready).toHaveBeenCalled();
      expect(manager.getBackend()).toBe('webgl');
    });

    it('should fall back to cpu on backend failure', async () => {
      (tf.setBackend as any).mockRejectedValueOnce(new Error('WebGL not supported'));
      
      await manager.initialize('webgl');
      
      expect(tf.setBackend).toHaveBeenCalledWith('cpu');
      expect(manager.getBackend()).toBe('cpu');
    });

    it('should not reinitialize', async () => {
      await manager.initialize('webgl');
      await manager.initialize('cpu'); // Should be ignored
      
      expect(tf.setBackend).toHaveBeenCalledTimes(1);
    });
  });

  describe('loadModel', () => {
    const mockModel = {
      predict: vi.fn((input) => ({
        dispose: vi.fn(),
        data: vi.fn().mockResolvedValue(new Float32Array([0.5, 0.5])),
      })),
      inputs: [{ shape: [null, 1024] }],
      outputs: [{ shape: [null, 360] }],
      dispose: vi.fn(),
    };

    it('should load LayersModel', async () => {
      (tf.loadLayersModel as any).mockResolvedValueOnce(mockModel);
      
      await manager.loadModel('/models/test.json');
      
      expect(tf.loadLayersModel).toHaveBeenCalledWith('/models/test.json');
      expect(manager.isLoaded()).toBe(true);
    });

    it('should fall back to GraphModel', async () => {
      (tf.loadLayersModel as any).mockRejectedValueOnce(new Error('Not a LayersModel'));
      (tf.loadGraphModel as any).mockResolvedValueOnce(mockModel);
      
      await manager.loadModel('/models/test.json');
      
      expect(tf.loadGraphModel).toHaveBeenCalledWith('/models/test.json');
      expect(manager.isLoaded()).toBe(true);
    });

    it('should throw if both model types fail', async () => {
      (tf.loadLayersModel as any).mockRejectedValueOnce(new Error('Failed'));
      (tf.loadGraphModel as any).mockRejectedValueOnce(new Error('Failed'));
      
      await expect(manager.loadModel('/models/test.json')).rejects.toThrow('Failed to load model');
    });

    it('should auto-initialize if needed', async () => {
      (tf.loadLayersModel as any).mockResolvedValueOnce(mockModel);
      
      await manager.loadModel('/models/test.json', { backend: 'cpu' });
      
      expect(tf.setBackend).toHaveBeenCalledWith('cpu');
    });
  });

  describe('predict', () => {
    const mockModel = {
      predict: vi.fn((input) => ({
        dispose: vi.fn(),
        data: vi.fn().mockResolvedValue(new Float32Array([0.5, 0.5])),
      })),
      inputs: [{ shape: [null, 1024] }],
      outputs: [{ shape: [null, 360] }],
      dispose: vi.fn(),
    };

    beforeEach(async () => {
      (tf.loadLayersModel as any).mockResolvedValueOnce(mockModel);
      await manager.initialize();
      await manager.loadModel('/models/test.json', { warmup: false });
    });

    it('should run prediction', async () => {
      const input = new Float32Array([1, 2, 3]);
      
      await manager.predict(input);
      
      expect(mockModel.predict).toHaveBeenCalled();
    });

    it('should throw if model not loaded', async () => {
      const unloadedManager = new TFJSModelManager();
      const input = new Float32Array([1, 2, 3]);
      
      await expect(unloadedManager.predict(input)).rejects.toThrow('Model not loaded');
    });

    it('should accept tensor input', async () => {
      const tensor = tf.tensor([1, 2, 3]);
      
      await manager.predict(tensor);
      
      expect(mockModel.predict).toHaveBeenCalled();
    });
  });

  describe('predictArray', () => {
    const mockModel = {
      predict: vi.fn((input) => ({
        dispose: vi.fn(),
        data: vi.fn().mockResolvedValue(new Float32Array([0.5, 0.5])),
      })),
      inputs: [{ shape: [null, 1024] }],
      outputs: [{ shape: [null, 360] }],
      dispose: vi.fn(),
    };

    beforeEach(async () => {
      (tf.loadLayersModel as any).mockResolvedValueOnce(mockModel);
      await manager.initialize();
      await manager.loadModel('/models/test.json', { warmup: false });
    });

    it('should return Float32Array', async () => {
      const input = new Float32Array([1, 2, 3]);
      const result = await manager.predictArray(input);
      
      expect(result).toBeInstanceOf(Float32Array);
    });
  });

  describe('model info', () => {
    const mockModel = {
      predict: vi.fn(),
      inputs: [{ shape: [null, 1024] }],
      outputs: [{ shape: [null, 360] }],
      dispose: vi.fn(),
    };

    beforeEach(async () => {
      (tf.loadLayersModel as any).mockResolvedValueOnce(mockModel);
      await manager.initialize();
      await manager.loadModel('/models/test.json', { warmup: false });
    });

    it('should get input shape', () => {
      const shape = manager.getInputShape();
      expect(shape).toEqual([null, 1024]);
    });

    it('should get output shape', () => {
      const shape = manager.getOutputShape();
      expect(shape).toEqual([null, 360]);
    });

    it('should return null if model not loaded', () => {
      const unloadedManager = new TFJSModelManager();
      expect(unloadedManager.getInputShape()).toBeNull();
      expect(unloadedManager.getOutputShape()).toBeNull();
    });
  });

  describe('dispose', () => {
    const mockModel = {
      predict: vi.fn(),
      inputs: [{ shape: [null, 1024] }],
      outputs: [{ shape: [null, 360] }],
      dispose: vi.fn(),
    };

    it('should dispose model', async () => {
      (tf.loadLayersModel as any).mockResolvedValueOnce(mockModel);
      await manager.initialize();
      await manager.loadModel('/models/test.json', { warmup: false });
      
      manager.dispose();
      
      expect(mockModel.dispose).toHaveBeenCalled();
      expect(manager.isLoaded()).toBe(false);
    });
  });

  describe('getMemoryInfo', () => {
    it('should return memory info', () => {
      const info = manager.getMemoryInfo();
      expect(info).toHaveProperty('numTensors');
      expect(info).toHaveProperty('numDataBuffers');
    });
  });
});

describe('TensorUtils', () => {
  describe('normalizeTensor', () => {
    it('should normalize tensor', () => {
      const tensor = tf.tensor([1, 2, 3]);
      const normalized = TensorUtils.normalizeTensor(tensor);
      expect(normalized).toBeDefined();
    });
  });

  describe('meanNormalize', () => {
    it('should apply mean normalization', () => {
      const tensor = tf.tensor([1, 2, 3]);
      const normalized = TensorUtils.meanNormalize(tensor);
      expect(normalized).toBeDefined();
    });
  });

  describe('toBatchedTensor', () => {
    it('should create batched tensor', () => {
      const data = new Float32Array([1, 2, 3, 4]);
      const batched = TensorUtils.toBatchedTensor(data, 2);
      expect(batched).toBeDefined();
    });
  });

  describe('argmax', () => {
    it('should get argmax', async () => {
      const tensor = tf.tensor([0.1, 0.9, 0.3]);
      const idx = await TensorUtils.argmax(tensor);
      expect(typeof idx).toBe('number');
    });
  });

  describe('topK', () => {
    it('should get top-k values', async () => {
      const tensor = tf.tensor([0.1, 0.9, 0.3, 0.7, 0.5]);
      const result = await TensorUtils.topK(tensor, 3);
      
      expect(result.indices).toHaveLength(3);
      expect(result.values).toHaveLength(3);
    });
  });

  describe('disposeTensors', () => {
    it('should dispose multiple tensors', () => {
      const t1 = tf.tensor([1, 2]);
      const t2 = tf.tensor([3, 4]);
      
      TensorUtils.disposeTensors(t1, t2);
      
      expect(t1.dispose).toHaveBeenCalled();
      expect(t2.dispose).toHaveBeenCalled();
    });

    it('should handle null tensors', () => {
      expect(() => TensorUtils.disposeTensors(null, undefined)).not.toThrow();
    });
  });
});

describe('createModelManager', () => {
  it('should create and initialize model manager', async () => {
    const mockModel = {
      predict: vi.fn(),
      inputs: [{ shape: [null, 1024] }],
      outputs: [{ shape: [null, 360] }],
      dispose: vi.fn(),
    };

    (tf.loadLayersModel as any).mockResolvedValueOnce(mockModel);
    
    const manager = await createModelManager('/models/test.json', { backend: 'cpu' });
    
    expect(manager.isLoaded()).toBe(true);
    expect(manager.getBackend()).toBe('cpu');
  });
});