import ctypes
import os
from ctypes import *

# ==========================================
# 1. 结构体定义
# ==========================================

class Engine_Identity(Structure):
    _fields_ = [
        ("Identity_Ok", c_int),
        ("Cof_Eff_Low_Identity", c_double),
        ("Cof_Eff_High_Identity", c_double)
    ]

class EngInPut(Structure):
    _fields_ = [
        ("Altp", c_double),           # 飞行高度 m
        ("Ma0", c_double),            # 马赫数
        ("dT0", c_double),            # 温差 K
        ("StepTime", c_double),       # 步长
        ("Wf", c_double),             # 主燃油流量 kg/h
        ("Wf_After", c_double),       # 加力燃油  kg/h
        ("Angle_FanVane", c_double),  # 风扇导叶
        ("Angle_CompVane", c_double), # 高压导叶
        ("A8", c_double),             # 喉道面积 m^2
        ("AddPower", c_double)        # 起动力矩 W
    ]

class EngOutPut(Structure):
    _fields_ = [
        ("NL", c_double),             # 风扇转速                (0,1)
        ("NH", c_double),             # 高压转速                (0,1)
        ("T1t", c_double),            # 进口温度                K
        ("P1t", c_double),            # 进口压力                kPa
        ("P1s", c_double),            # 进口静压                kPa
        ("P3s", c_double),            # 压气机出口压力          kPa
        ("T5t", c_double),            # 涡轮出口温度            K
        ("P5t", c_double),            # 涡轮出口压力            kPa
        ("Wf_Main", c_double),        # 实际主燃油流量          kg/h
        ("T5t_Gas", c_double)         # 涡轮出口实际气体温度     K
    ]

# ==========================================
# 2. DLL 封装类
# ==========================================

class AeroEngineDLL:
    def __init__(self, dll_path=None):
        if dll_path is None:
            dll_path = os.path.join(os.path.dirname(__file__), "libEngine.dll")
        
        self.lib = WinDLL(dll_path) # 使用 WinDLL 因为是 stdcall

        # 配置函数原型
        self.lib.CreateEng.argtypes = [POINTER(POINTER(c_int)), Engine_Identity]
        self.lib.CreateEng.restype = EngOutPut

        self.lib.EngStepGo.argtypes = [POINTER(c_int), EngInPut]
        self.lib.EngStepGo.restype = EngOutPut

        self.lib.DestroyEng.argtypes = [POINTER(c_int)]
        self.lib.DestroyEng.restype = EngOutPut

        self.h_engine = None
        self.dt = 0.02 # 默认仿真步长 20ms

    def reset(self):
        # 如果已经存在实例，先销毁
        if self.h_engine:
            self.lib.DestroyEng(self.h_engine)
        
        # 创建新实例
        self.h_engine = POINTER(c_int)()
        identity = Engine_Identity(1, 1.0, 1.0)
        self.lib.CreateEng(byref(self.h_engine), identity)
        
        # 初始化一个默认的输入状态
        self.current_input = EngInPut()
        self.current_input.Altp = 0.0
        self.current_input.Ma0 = 0.0
        self.current_input.dT0 = 0.0
        self.current_input.StepTime = self.dt
        self.current_input.Wf = 100.0  # 初始燃油
        self.current_input.A8 = 0.1    # 初始 A8
        self.current_input.AddPower = 0.0
        
        # 跑一步获取初始状态
        out = self.lib.EngStepGo(self.h_engine, self.current_input)
        return out

    def step(self, action_dict):
        """
        action_dict: 包含具体物理值的字典
        {'Wf': 200.0, 'A8': 0.12, ...}
        """
        # 更新输入结构体
        if 'Wf' in action_dict: self.current_input.Wf = action_dict['Wf']
        if 'A8' in action_dict: self.current_input.A8 = action_dict['A8']
        if 'FanVane' in action_dict: self.current_input.Angle_FanVane = action_dict['FanVane']
        if 'CompVane' in action_dict: self.current_input.Angle_CompVane = action_dict['CompVane']
        if 'AddPower' in action_dict: self.current_input.AddPower = action_dict['AddPower']
        
        # 确保步长正确
        self.current_input.StepTime = self.dt
        
        # 调用 DLL
        out = self.lib.EngStepGo(self.h_engine, self.current_input)
        return out

    def close(self):
        if self.h_engine:
            try:
                self.lib.DestroyEng(self.h_engine)
            except OSError:
                # 忽略销毁时的访问冲突，可能是 DLL 内部状态问题
                pass
            self.h_engine = None
        
        # 尝试释放 DLL
        # 注意：这通常不是必须的，但在某些情况下（如反复加载/卸载或 DLL 内部有全局状态）可能有帮助
        if hasattr(self, 'lib'):
            try:
                import _ctypes
                _ctypes.FreeLibrary(self.lib._handle)
            except:
                pass
            del self.lib