import 'dart:typed_data';

import 'package:tflite_flutter/tflite_flutter.dart';

import 'detection_result.dart';
import 'nms.dart';

/// YOLOv8n 端侧推理实现（对应 Kotlin TFLiteDetector）。
/// 模型输出布局（ultralytics litert 导出）：[1, 4 + nc, anchors]，
/// cx/cy/w/h 已归一化，类别得分已过 sigmoid；按 out[dim][anchor] 索引。
/// 输入为 NCHW [1, 3, 704, 704]（litert 导出保留 torch 布局）。
class TfliteDetector {
  static const int inputSize = 704;
  // 野鸡数据置信度普遍偏低（0.1~0.2 量级），保留低分池供运动检测提升
  static const double minScore = 0.10;
  static const double iouThreshold = 0.45;
  static const int maxDetections = 20;
  static const String modelAsset = 'assets/model.tflite';
  static const String labelsAsset = 'assets/labels.txt';

  final Interpreter _interpreter;
  final List<String> _labels;
  final int _numClasses;
  final int _numAnchors;

  final Float32List _input =
      Float32List(1 * inputSize * inputSize * 3);

  /// 输出按模型形状 [1, 4+nc, anchors] 的嵌套 List 组织，
  /// run() 要求输出对象形状与模型完全一致（扁平 List 会被拒）。
  final List<List<List<double>>> _output;

  TfliteDetector._(this._interpreter, this._labels, this._numClasses,
      this._numAnchors, this._output);

  /// 模型缺失或加载失败返回 null（App 降级为仅预览）。
  /// 在后台 isolate 内调用（模型字节由主 isolate 读取后传入）。
  static Future<TfliteDetector?> fromBuffer(
      Uint8List bytes, List<String> labels) async {
    try {
      final interpreter = Interpreter.fromBuffer(
        bytes,
        options: InterpreterOptions()..threads = 4,
      );
      return TfliteDetector._fromModel(interpreter, labels);
    } catch (_) {
      return null;
    }
  }

  /// 输出布局 [1, 4+nc, anchors] 取自模型本身，类别数不与 labels 文件长度耦合。
  factory TfliteDetector._fromModel(
      Interpreter interpreter, List<String> labels) {
    final shape = interpreter.getOutputTensor(0).shape;
    final numClasses =
        shape.length >= 3 && shape[1] > 4 ? shape[1] - 4 : labels.length;
    final numAnchors = shape.length >= 3 && shape[2] > 0 ? shape[2] : 2100;
    final output = List.generate(
      1,
      (_) => List.generate(
        numClasses + 4,
        (_) => List<double>.filled(numAnchors, 0),
      ),
    );
    return TfliteDetector._(
        interpreter, labels, numClasses, numAnchors, output);
  }

  /// 原始数据接口（后台 isolate 用，不依赖 CameraImage）。
  /// 输出坐标统一反算为原图归一化空间（与 MotionDetector 一致），
  /// 否则 CENTER_CROP 裁剪偏移会让检测框系统性偏移。
  List<DetectionResult> detectRaw({
    required List<Uint8List> planes,
    required List<int> strides,
    required int width,
    required int height,
    required bool isBgra,
  }) {
    preprocess(
        planes: planes,
        strides: strides,
        width: width,
        height: height,
        isBgra: isBgra);
    // 传原始字节视图而非 Float32List：tflite_flutter 会对非 ByteBuffer/Uint8List
    // 输入调用 resizeInputTensor（1 维 [1486848]），使 node 0 TRANSPOSE prepare 失败
    _interpreter.run(_input.buffer.asUint8List(), _output);
    final dets = postprocess();
    // 反算与 preprocess 的 scale/dx/dy 公式一致（704 输入空间 → 原图归一化）
    final scale = inputSize / width < inputSize / height
        ? inputSize / width
        : inputSize / height;
    final dx = (inputSize - width * scale) / 2;
    final dy = (inputSize - height * scale) / 2;
    if (dx == 0 && dy == 0) return dets;
    return dets
        .map((r) => r.copyWith(
              left: (r.left * inputSize - dx) / (width * scale),
              right: (r.right * inputSize - dx) / (width * scale),
              top: (r.top * inputSize - dy) / (height * scale),
              bottom: (r.bottom * inputSize - dy) / (height * scale),
            ))
        .toList();
  }

  /// 按像素格式分派：iOS bgra8888 单平面 / Android yuv420 多平面。
  void preprocess({
    required List<Uint8List> planes,
    required List<int> strides,
    required int width,
    required int height,
    required bool isBgra,
  }) {
    if (isBgra) {
      _preprocessBgra(planes[0], strides[0], width, height);
    } else {
      _preprocessYuv(planes, strides, width, height);
    }
  }

  /// BGRA8888 单平面（iOS）：每像素 4 字节 [b,g,r,a]，双线性采样，
  /// letterbox（等比缩到长边 704，短边黑边补 0，与 YOLO 训练一致）。
  void _preprocessBgra(Uint8List src, int stride, int srcW, int srcH) {
    final plane = inputSize * inputSize;
    final scale = inputSize / srcW < inputSize / srcH
        ? inputSize / srcW
        : inputSize / srcH;
    final dx = (inputSize - srcW * scale) / 2;
    final dy = (inputSize - srcH * scale) / 2;

    for (var oy = 0; oy < inputSize; oy++) {
      final syf = (oy - dy) / scale;
      if (syf < 0 || syf >= srcH) {
        for (var ox = 0; ox < inputSize; ox++) {
          final p = oy * inputSize + ox;
          _input[p] = 0;
          _input[p + plane] = 0;
          _input[p + 2 * plane] = 0;
        }
        continue;
      }
      for (var ox = 0; ox < inputSize; ox++) {
        final p = oy * inputSize + ox;
        final sxf = (ox - dx) / scale;
        if (sxf < 0 || sxf >= srcW) {
          _input[p] = 0;
          _input[p + plane] = 0;
          _input[p + 2 * plane] = 0;
          continue;
        }
        final x0 = sxf.floor(), y0 = syf.floor();
        final x1 = x0 < srcW - 1 ? x0 + 1 : x0;
        final y1 = y0 < srcH - 1 ? y0 + 1 : y0;
        final fx = sxf - x0, fy = syf - y0;

        // BGRA 字节序：+0 B、+1 G、+2 R、+3 A
        final i00 = y0 * stride + x0 * 4;
        final i10 = y0 * stride + x1 * 4;
        final i01 = y1 * stride + x0 * 4;
        final i11 = y1 * stride + x1 * 4;
        final r00 = src[i00 + 2].toDouble();
        final g00 = src[i00 + 1].toDouble();
        final b00 = src[i00].toDouble();
        final r10 = src[i10 + 2].toDouble();
        final g10 = src[i10 + 1].toDouble();
        final b10 = src[i10].toDouble();
        final r01 = src[i01 + 2].toDouble();
        final g01 = src[i01 + 1].toDouble();
        final b01 = src[i01].toDouble();
        final r11 = src[i11 + 2].toDouble();
        final g11 = src[i11 + 1].toDouble();
        final b11 = src[i11].toDouble();

        _input[p] = _bl(r00, r10, r01, r11, fx, fy) / 255.0;
        _input[p + plane] = _bl(g00, g10, g01, g11, fx, fy) / 255.0;
        _input[p + 2 * plane] = _bl(b00, b10, b01, b11, fx, fy) / 255.0;
      }
    }
  }

  /// letterbox 缩放 + YUV → RGB 归一化 0~1（NCHW），双线性采样。
  /// 兼容 NV12（iOS 双平面，UV 交错）与 I420（Android 三平面）。
  void _preprocessYuv(
      List<Uint8List> planes, List<int> strides, int srcW, int srcH) {
    final plane = inputSize * inputSize;
    final y = planes[0];
    final nv12 = planes.length == 2;
    final uv = nv12 ? planes[1] : null;
    final u = nv12 ? null : planes[1];
    final v = nv12 ? null : planes[2];
    final yStride = strides[0];
    final uvStride = strides[1];

    // U/V 平面采样（nv12：偶位 U 奇位 V；i420：三平面分离）
    double uAt(int x, int y) => nv12
        ? uv![y * uvStride + x * 2] - 128.0
        : u![y * uvStride + x] - 128.0;
    double vAt(int x, int y) => nv12
        ? uv![y * uvStride + x * 2 + 1] - 128.0
        : v![y * uvStride + x] - 128.0;

    final scale = inputSize / srcW < inputSize / srcH
        ? inputSize / srcW
        : inputSize / srcH;
    final dx = (inputSize - srcW * scale) / 2;
    final dy = (inputSize - srcH * scale) / 2;

    for (var oy = 0; oy < inputSize; oy++) {
      final syf = (oy - dy) / scale;
      if (syf < 0 || syf >= srcH) {
        for (var ox = 0; ox < inputSize; ox++) {
          final p = oy * inputSize + ox;
          _input[p] = 0;
          _input[p + plane] = 0;
          _input[p + 2 * plane] = 0;
        }
        continue;
      }
      for (var ox = 0; ox < inputSize; ox++) {
        final p = oy * inputSize + ox;
        final sxf = (ox - dx) / scale;
        if (sxf < 0 || sxf >= srcW) {
          _input[p] = 0;
          _input[p + plane] = 0;
          _input[p + 2 * plane] = 0;
          continue;
        }
        final x0 = sxf.floor(), y0 = syf.floor();
        final x1 = x0 < srcW - 1 ? x0 + 1 : x0;
        final y1 = y0 < srcH - 1 ? y0 + 1 : y0;
        final fx = sxf - x0, fy = syf - y0;

        // Y 双线性
        final y00 = y[y0 * yStride + x0].toDouble();
        final y10 = y[y0 * yStride + x1].toDouble();
        final y01 = y[y1 * yStride + x0].toDouble();
        final y11 = y[y1 * yStride + x1].toDouble();
        final yy = _bl(y00, y10, y01, y11, fx, fy);

        // U/V 双线性（4:2:0 半分辨率，按像素坐标定位后除 2）
        final maxUx = srcW ~/ 2 - 1;
        final maxUy = srcH ~/ 2 - 1;
        final ux0 = (x0 ~/ 2).clamp(0, maxUx).toInt();
        final uy0 = (y0 ~/ 2).clamp(0, maxUy).toInt();
        final ux1 = (x1 ~/ 2).clamp(0, maxUx).toInt();
        final uy1 = (y1 ~/ 2).clamp(0, maxUy).toInt();
        final u00 = uAt(ux0, uy0);
        final u10 = uAt(ux1, uy0);
        final u01 = uAt(ux0, uy1);
        final u11 = uAt(ux1, uy1);
        final uu = _bl(u00, u10, u01, u11, fx, fy);

        final v00 = vAt(ux0, uy0);
        final v10 = vAt(ux1, uy0);
        final v01 = vAt(ux0, uy1);
        final v11 = vAt(ux1, uy1);
        final vv = _bl(v00, v10, v01, v11, fx, fy);

        // 有限范围展开（VideoRange Y 16~235，Cb/Cr 16~240）
        final yr = (yy - 16.0) * (255.0 / 219.0);
        final un = uu * (255.0 / 224.0);
        final vn = vv * (255.0 / 224.0);

        // NCHW：r/g/b 分平面存储
        _input[p] = (yr + 1.402 * vn) / 255.0;
        _input[p + plane] = (yr - 0.344136 * un - 0.714136 * vn) / 255.0;
        _input[p + 2 * plane] = (yr + 1.772 * un) / 255.0;
      }
    }
  }

  static double _bl(double a, double b, double c, double d, double fx,
          double fy) =>
      (1 - fx) * (1 - fy) * a + fx * (1 - fy) * b +
      (1 - fx) * fy * c + fx * fy * d;

  List<DetectionResult> postprocess() {
    final out = _output[0];
    final boxes = <DetectionResult>[];
    for (var a = 0; a < _numAnchors; a++) {
      final cx = out[0][a];
      final cy = out[1][a];
      final w = out[2][a];
      final h = out[3][a];
      var bestCls = 0;
      var bestScore = 0.0;
      for (var c = 0; c < _numClasses; c++) {
        final s = out[4 + c][a];
        if (s > bestScore) {
          bestScore = s;
          bestCls = c;
        }
      }
      final label =
          bestCls < _labels.length ? _labels[bestCls] : 'unknown';
      // 低分池保留，供运动检测提升显示
      if (bestScore < minScore) continue;
      boxes.add(DetectionResult(
        label: label,
        score: bestScore,
        left: (cx - w / 2).clamp(0.0, 1.0),
        top: (cy - h / 2).clamp(0.0, 1.0),
        right: (cx + w / 2).clamp(0.0, 1.0),
        bottom: (cy + h / 2).clamp(0.0, 1.0),
      ));
    }
    final kept = nms(boxes, iouThreshold);
    return kept.take(maxDetections).toList();
  }

  void dispose() => _interpreter.close();
}