using Godot;
using System;
using Microsoft.ML.OnnxRuntime;
using System.Linq;
using Microsoft.ML.OnnxRuntime.Tensors;
using System.Collections.Generic;
using System.IO;
using System.Text;
using OpenCvSharp;

public class MultiTaskPrediction
{
    public float[] Position { get; set; }  // [x, y] coordinates
    public float FlagProbability { get; set; }  // Probability of blink
}


public partial class NeuralNetwork : Node
{
	// Path to the encrypted ONNX model file
	[Export]
	public string EncryptedModelPath { get; set; } = "";

	public bool UsingDirectML { get; private set; }

	// ONNX inference session
	private InferenceSession _session;

	// Model input parameters
	private string _inputName;
	private string _eyeSideInputName;  // eye_side, 0 is left, 1 is right	
	private int _inputHeight = 256;
	private int _inputWidth = 256;
	private int _inputChannels = 1; // 1 for grayscale

	// Model output parameters
	private string _positionOutputName = "position";
	private string _flagOutputName = "flag_prob";

	// Called when the node enters the scene tree for the first time
	public override void _Ready()
	{
	}

	public void LoadModel()
	{
		GD.Print($"Loading model from: {EncryptedModelPath}");

		// Convert the Godot resource path to a system path
		string systemPath = ProjectSettings.GlobalizePath(EncryptedModelPath);

		// Load and decrypt the model
		byte[] modelBytes = DecryptModel(systemPath);

		// Create session options
		var sessionOptions = new SessionOptions();
		var availableProviders = OrtEnv.Instance().GetAvailableProviders();
		bool hasDirectML = availableProviders.Any(provider => provider == "DmlExecutionProvider");

		// Try to use DirectML if available, with proper exception handling
		if (hasDirectML)
		{
			try
			{
				GD.Print("Attempting to use DirectML acceleration");
				sessionOptions.AppendExecutionProvider_DML(0);
				UsingDirectML = true;

				GD.Print("DirectML provider added successfully");
			}
			catch (Exception ex)
			{
				GD.PrintErr($"Failed to initialize DirectML: {ex.Message}");
				GD.Print("Falling back to CPU only");
				UsingDirectML = false;

				// Clear any providers that might have been added
				sessionOptions = new SessionOptions();
			}
		}
		else
		{
			GD.Print("DirectML not available, using CPU");
			UsingDirectML = false;
        }

		// Always add CPU as fallback
		sessionOptions.AppendExecutionProvider_CPU();

		// Load the model from decrypted bytes
		_session = new InferenceSession(modelBytes, sessionOptions);

		// Get model input details
		var inputMeta = _session.InputMetadata.First();
		_inputName = inputMeta.Key;
		_eyeSideInputName = _session.InputMetadata.Keys.ElementAt(1);
		var shape = inputMeta.Value.Dimensions;

		// Determine input dimensions
		_inputChannels = shape.Length > 1 ? (int)shape[1] : 1;
		_inputHeight = shape.Length > 2 ? (int)shape[2] : 256;
		_inputWidth = shape.Length > 3 ? (int)shape[3] : 256;

		// Get model output names
		var outputNames = _session.OutputMetadata.Keys.ToArray();
		if (outputNames.Length >= 2)
		{
			_positionOutputName = outputNames[0];
			_flagOutputName = outputNames[1];
		}

		// Log model information
		GD.Print($"Model loaded successfully");
		GD.Print($"Input name: {_inputName}");
		GD.Print($"Eye side input name: {_eyeSideInputName}");
		GD.Print($"Input shape: 1x{_inputChannels}x{_inputHeight}x{_inputWidth}");
		GD.Print($"Position output name: {_positionOutputName}");
		GD.Print($"Flag output name: {_flagOutputName}");

		// Log model outputs
		GD.Print("Model outputs:");
		foreach (var output in _session.OutputMetadata)
		{
			GD.Print($"  - {output.Key}: {string.Join("x", output.Value.Dimensions)}");
		}
	}

	private byte[] DecryptModel(string obfuscatedModelPath)
	{
		byte[] part1 = Encoding.UTF8.GetBytes("gFlFTSp-sW-");
		byte[] part2 = Encoding.UTF8.GetBytes("HI");
		byte[] part3 = Encoding.UTF8.GetBytes("wuQvfbLcI1u6wsRc2b");
		byte[] part4 = Encoding.UTF8.GetBytes("3-eOb5SpPc_Y=");

		byte[] key = new byte[part1.Length + part2.Length + part3.Length + part4.Length];
		Buffer.BlockCopy(part1, 0, key, 0, part1.Length);
		Buffer.BlockCopy(part2, 0, key, part1.Length, part2.Length);
		Buffer.BlockCopy(part3, 0, key, part1.Length + part2.Length, part3.Length);
		Buffer.BlockCopy(part4, 0, key, part1.Length + part2.Length + part3.Length, part4.Length);

		byte[] data = File.ReadAllBytes(obfuscatedModelPath);
		byte[] result = new byte[data.Length];

		for (int i = 0; i < data.Length; i++)
			result[i] = (byte)(data[i] ^ key[i % key.Length]);

		return result;
	}

	private void SaveDebugModel(byte[] data, string filename)
	{
		string debugDir = "debug";
		if (!Directory.Exists(debugDir))
		{
			Directory.CreateDirectory(debugDir);
		}
		string debugFilePath = Path.Combine(debugDir, filename);
		File.WriteAllBytes(debugFilePath, data);
		GD.Print($"Saved debug file: {debugFilePath} ({data.Length} bytes)");
	}

	public MultiTaskPrediction ProcessFrame(Mat frame, bool isRightEye = false)
	{
		// Preprocess the image
		var inputTensor = PreprocessFrame(frame);

		// Create eye side tensor (0 for left eye, 1 for right eye)
		var eyeSideTensor = new DenseTensor<float>(new[] { 1 });
		eyeSideTensor[0] = isRightEye ? 1.0f : 0.0f;

		// Run inference with both inputs
		var inputs = new List<NamedOnnxValue> 
		{ 
			NamedOnnxValue.CreateFromTensor(_inputName, inputTensor),
			NamedOnnxValue.CreateFromTensor(_eyeSideInputName, eyeSideTensor)
		};

		if (_session == null)
		{
			return null;
		}

        using var results = _session.Run(inputs);
        
        // Process the multi-task outputs
        var prediction = new MultiTaskPrediction();
		
		// If there is only one result, this is an old model without blink flags
		if (results.Count() == 1)
		{
			prediction.FlagProbability = 0.0f; // Default to no blink
			prediction.Position = results.First().AsTensor<float>().ToArray();
			return prediction;
		}
        
        // Get position output
		if (results.Any(r => r.Name == _positionOutputName))
		{
			var positionOutput = results.First(r => r.Name == _positionOutputName);
			if (positionOutput.Value is DenseTensor<float> positionTensor)
			{
				prediction.Position = positionTensor.ToArray();
			}
			else
			{
				GD.PrintErr("Position output is not a float tensor");
				return null;
			}
		}
		else
		{
			GD.PrintErr($"Position output '{_positionOutputName}' not found");
			return null;
		}
        
        // Get flag output
		if (results.Any(r => r.Name == _flagOutputName))
		{
			var flagOutput = results.First(r => r.Name == _flagOutputName);
			if (flagOutput.Value is DenseTensor<float> flagTensor)
			{
				prediction.FlagProbability = flagTensor[0]; // Single value
			}
			else
			{
				GD.PrintErr("Flag output is not a float tensor");
				return null;
			}
		}
		else
		{
			GD.PrintErr($"Flag output '{_flagOutputName}' not found");
			return null;
		}
        
        return prediction;
    }

	private DenseTensor<float> PreprocessFrame(Mat frame)
	{
		// Clone the input to avoid modifying the original
		using (Mat processedFrame = frame.Clone())
		{
			// Convert to the correct format if needed
			if (_inputChannels == 1 && processedFrame.Channels() > 1)
			{
				Cv2.CvtColor(processedFrame, processedFrame, ColorConversionCodes.BGR2GRAY);
			}

			// Resize if needed
			if (processedFrame.Width != _inputWidth || processedFrame.Height != _inputHeight)
			{
				Cv2.Resize(processedFrame, processedFrame, new Size(_inputWidth, _inputHeight), 0, 0, InterpolationFlags.Linear);
			}

			// Prepare the input tensor (NCHW format)
			var inputTensor = new DenseTensor<float>(new[] { 1, _inputChannels, _inputHeight, _inputWidth });

			// Normalize and copy image data to tensor
			if (_inputChannels == 1)
			{
				// Grayscale image (1 channel)
				for (int h = 0; h < _inputHeight; h++)
				{
					for (int w = 0; w < _inputWidth; w++)
					{
						float pixelValue = processedFrame.At<byte>(h, w) / 255.0f;
						inputTensor[0, 0, h, w] = pixelValue;
					}
				}
			}

			return inputTensor;
		}
	}

	public bool IsFlagCondition(MultiTaskPrediction prediction, float threshold = 0.5f)
	{
		return prediction?.FlagProbability > threshold;
	}
}