using System.Collections.Generic;
using System.Threading.Tasks;
using Unity.InferenceEngine;
using UnityEngine;
using SentisModels;

namespace OnDeviceAgent.Inference
{
public abstract class YoloDetector : MonoBehaviour
{
    [SerializeField, Tooltip("Optional StreamingAssets-relative model override; blank uses the subclass default.")]
    string m_ModelPath = string.Empty;
    [SerializeField] BackendType m_BackendType = BackendType.GPUCompute;
    [SerializeField, Range(0.05f, 0.95f)] float m_ConfidenceThreshold = 0.25f;
    [SerializeField, Tooltip("Log the raw model output shapes on the first inference (diagnostic).")]
    bool m_LogOutputShapeOnce = true;
    [SerializeField, Tooltip("Run one dummy inference when the model loads so the GPUCompute shaders " +
        "compile up front; otherwise the first real detection stalls while every kernel is compiled.")]
    bool m_WarmupOnEnable = true;
    [SerializeField, Tooltip("Swap R<->B before inference. YOLOX is typically trained on BGR while Sentis " +
        "ToTensor produces RGB. If detections are wrong (false positives for unrelated classes), toggle this.")]
    bool m_SwapRedBlue = true;

    protected abstract string DefaultModelPath { get; }

    YoloxDetector m_Core;

    public struct Detection
    {
        public int ClassId;
        public string ClassName;
        public float Confidence;
        public Rect BoxXyxy;
    }

    public bool IsReady => m_Core != null && m_Core.IsReady;

    void OnEnable()
    {
        LoadModel();
    }

    void OnDisable()
    {
        m_Core?.Dispose();
        m_Core = null;
    }

    void LoadModel()
    {
        m_Core?.Dispose();
        m_Core = new YoloxDetector(
            m_BackendType,
            m_ConfidenceThreshold,
            m_SwapRedBlue,
            m_WarmupOnEnable,
            m_LogOutputShapeOnce);

        try
        {
            var root = ModelRootProvisioner.EnsureModelRoot("Model/YOLO");
            m_Core.Load(root, ModelFileRelativeToRoot());
        }
        catch (System.Exception exception)
        {
            Debug.LogException(exception);
        }
    }

    string GetModelPath()
    {
        return string.IsNullOrWhiteSpace(m_ModelPath) ? DefaultModelPath : m_ModelPath.Trim();
    }

    // Strips the "Model/YOLO/" package-root prefix from the configured path to get the file name the package
    // Load expects relative to its modelRoot.
    string ModelFileRelativeToRoot()
    {
        const string prefix = "Model/YOLO/";
        var path = GetModelPath().Replace('\\', '/');
        return path.StartsWith(prefix, System.StringComparison.Ordinal) ? path.Substring(prefix.Length) : path;
    }

    public async Task<List<Detection>> DetectAsync(Texture frame)
    {
        if (m_Core == null)
            return new List<Detection>();

        var raw = await m_Core.DetectAsync(frame);
        var result = new List<Detection>(raw.Count);
        foreach (var d in raw)
            result.Add(new Detection { ClassId = d.ClassId, ClassName = d.ClassName, Confidence = d.Confidence, BoxXyxy = d.BoxXyxy });
        return result;
    }
}
}