import cv from '@techstark/opencv-js'

export interface TrackerParams {
  lowerYellow: any;
  upperYellow: any;
  blurSize: number;
  morphKernel: number;
  minContourArea: number;
  useRoi: boolean;
  roiScaleX: number;
  roiScaleY: number;
}

export interface Position {
  x: number;
  y: number;
}

export class YellowBallTracker {
  public srcMat: cv.Mat | null = null
  private hsvMat: cv.Mat | null = null
  private blurredMat: cv.Mat | null = null
  private maskMat: cv.Mat | null = null
  private hierarchy: cv.Mat | null = null
  private contours: cv.MatVector | null = null
  private morphKernelMat: cv.Mat | null = null
  private lowerBoundMat: cv.Mat | null = null
  private upperBoundMat: cv.Mat | null = null
  private params: TrackerParams

  constructor(params: TrackerParams) {
    this.params = params
  }

  public updateKernel() {
    if (this.morphKernelMat) this.morphKernelMat.delete()
    this.morphKernelMat = cv.Mat.ones(this.params.morphKernel, this.params.morphKernel, cv.CV_8U)
  }

  public initMats(w: number, h: number) {
    this.cleanupMats()
    
    this.srcMat = new cv.Mat(h, w, cv.CV_8UC4)
    this.hsvMat = new cv.Mat()
    this.blurredMat = new cv.Mat()
    this.maskMat = new cv.Mat()
    this.hierarchy = new cv.Mat()
    this.contours = new cv.MatVector()
    this.updateKernel()
    // 创建阈值 Mat
    this.lowerBoundMat = new cv.Mat(h, w, cv.CV_8UC3, this.params.lowerYellow)
    this.upperBoundMat = new cv.Mat(h, w, cv.CV_8UC3, this.params.upperYellow)
  }

  private calculateROI(lastPos: Position): { x: number, y: number, width: number, height: number } {
    const width = this.srcMat!.cols * this.params.roiScaleX
    const height = this.srcMat!.rows * this.params.roiScaleY
    const x = Math.max(0, Math.min(this.srcMat!.cols - width, lastPos.x - width / 2))
    const y = Math.max(0, Math.min(this.srcMat!.rows - height, lastPos.y - height / 2))
    
    return { 
      x: Math.round(x), 
      y: Math.round(y), 
      width: Math.round(width), 
      height: Math.round(height) 
    }
  }

  private drawROIDebugInfo(roi: { x: number, y: number, width: number, height: number }) {
    cv.rectangle(this.srcMat!,
      new cv.Point(roi.x, roi.y),
      new cv.Point(roi.x + roi.width, roi.y + roi.height),
      new cv.Scalar(255, 0, 0, 255), 2)
  }

  private processImage(inputMat: cv.Mat, lowerBound: cv.Mat, upperBound: cv.Mat) {
    // 颜色空间转换和模糊处理
    cv.cvtColor(inputMat, this.hsvMat!, cv.COLOR_RGBA2RGB)
    cv.cvtColor(this.hsvMat!, this.hsvMat!, cv.COLOR_RGB2HSV)
    cv.blur(this.hsvMat!, this.blurredMat!, new cv.Size(this.params.blurSize, this.params.blurSize))
    
    // 颜色范围过滤
    cv.inRange(this.blurredMat!, lowerBound, upperBound, this.maskMat!)
    
    // 形态学操作
    cv.morphologyEx(this.maskMat!, this.maskMat!, cv.MORPH_OPEN, this.morphKernelMat!)
    cv.morphologyEx(this.maskMat!, this.maskMat!, cv.MORPH_CLOSE, this.morphKernelMat!)
    cv.dilate(this.maskMat!, this.maskMat!, this.morphKernelMat!)
  }

  private findBestContour(offsetX: number = 0, offsetY: number = 0): Position | null {
    // 查找轮廓
    this.contours!.delete()
    this.contours = new cv.MatVector()
    cv.findContours(this.maskMat!, this.contours, this.hierarchy!, cv.RETR_EXTERNAL, cv.CHAIN_APPROX_SIMPLE)

    // 找到最大面积的轮廓
    let maxArea = 0
    let best: Position | null = null
    
    for (let i = 0; i < this.contours.size(); i++) {
      const cnt = this.contours.get(i)
      const area = cv.contourArea(cnt)
      
      if (area > this.params.minContourArea && area > maxArea) {
        maxArea = area
        const circle = cv.minEnclosingCircle(cnt)
        best = {
          x: Math.round(circle.center.x + offsetX),
          y: Math.round(circle.center.y + offsetY)
        }
      }
      cnt.delete()
    }
    
    return best
  }

  private drawDetectionResult(position: Position) {
    cv.circle(this.srcMat!, new cv.Point(position.x, position.y), 10, new cv.Scalar(0, 255, 0, 255), 2)
  }

  public detectYellowBall(imgData: ImageData, lastPos?: Position | null): Position | null {
    if (!this.srcMat || !this.hsvMat || !this.blurredMat || !this.maskMat || !this.hierarchy || 
        !this.contours || !this.morphKernelMat || !this.lowerBoundMat || !this.upperBoundMat) {
      console.error('Mats not initialized')
      return null
    }

    // Set image data
    this.srcMat.data.set(imgData.data)
    
    let best: Position | null = null
    let srcRoiMat = new cv.Mat()

    try {
      // 尝试ROI检测
      if (this.params.useRoi && lastPos) {
        const roi = this.calculateROI(lastPos)
        this.drawROIDebugInfo(roi)

        // 提取ROI区域
        const roiRect = new cv.Rect(roi.x, roi.y, roi.width, roi.height)
        srcRoiMat = this.srcMat.roi(roiRect)

        // 为ROI区域创建合适大小的阈值Mat
        const roiLowerBound = new cv.Mat(roi.height, roi.width, cv.CV_8UC3, this.params.lowerYellow)
        const roiUpperBound = new cv.Mat(roi.height, roi.width, cv.CV_8UC3, this.params.upperYellow)

        try {
          this.processImage(srcRoiMat, roiLowerBound, roiUpperBound)
          best = this.findBestContour(roi.x, roi.y)
        } finally {
          roiLowerBound.delete()
          roiUpperBound.delete()
        }
      }

      // 如果ROI中没找到或者没启用ROI，就在全图搜索
      if (!best) {
        if (this.params.useRoi && lastPos) {
          console.log("ROI失败，全域搜索")
        }

        this.processImage(this.srcMat, this.lowerBoundMat, this.upperBoundMat)
        best = this.findBestContour()
      }

      // 在原图上绘制找到的点
      if (best) {
        this.drawDetectionResult(best)
      }

    } finally {
      // 清理资源
      srcRoiMat.delete()
    }

    return best
  }

  public getOutputCanvas(targetCanvas: HTMLCanvasElement) {
    if (!this.srcMat) return null
    cv.imshow(targetCanvas, this.srcMat)
    return targetCanvas
  }

  public updateParams(params: Partial<TrackerParams>) {
    this.params = { ...this.params, ...params }
    
    // 如果更新了morphKernel，需要重新创建kernel
    if ('morphKernel' in params) {
      this.updateKernel()
    }
  }

  public cleanupMats() {
    // 清理资源
    if (this.srcMat) this.srcMat.delete()
    if (this.hsvMat) this.hsvMat.delete()
    if (this.blurredMat) this.blurredMat.delete()
    if (this.maskMat) this.maskMat.delete()
    if (this.hierarchy) this.hierarchy.delete()
    if (this.contours) this.contours.delete()
    if (this.morphKernelMat) this.morphKernelMat.delete()
    if (this.lowerBoundMat) this.lowerBoundMat.delete()
    if (this.upperBoundMat) this.upperBoundMat.delete()
    
    // 重置引用
    this.srcMat = null
    this.hsvMat = null
    this.blurredMat = null
    this.maskMat = null
    this.hierarchy = null
    this.contours = null
    this.morphKernelMat = null
    this.lowerBoundMat = null
    this.upperBoundMat = null
  }
}