-- EP-series piston engine driver for ArduPilot Lua scripting
--
-- Link layer (engine ECU <-> RS422 level shifter <-> free SERIALx, 115200 8N1):
--   ECU -> FC  : 0xAA 0x55 | 39 payload bytes | checksum byte (sum of payload & 0xFF)
--   FC  -> ECU : 0x55 0xAA | throttle(%x10, u16) | starter cmd | pad u16 | pad u8 | seq |
--                checksum byte (sum of bytes 3..9 & 0xFF)
--
-- Telemetry is forwarded to the ground station through the scripting EFI
-- backend, so EFI_TYPE must be set to 7 (SCRIPTING) and the engine serial
-- port to protocol 28.

local k_throttle_servo_fn = 70

local MAV_SEVERITY = {EMERGENCY=0, ALERT=1, CRITICAL=2, ERROR=3, WARNING=4, NOTICE=5, INFO=6, DEBUG=7}

-- ---------------------------------------------------------------- parameters

local PKEY = 64
local PPFX = "EP_"

-- name, index, default
local PARAM_DEFS = {
    { "ENABLE",      1, 0    },
    { "THR_FUNC",    2, 0    },
    { "THR_MIN",     3, 1600 },
    { "STR_FUNC",    4, 0    },
    { "STR_MIN",     5, 1800 },
    { "THR_PWM_MIN", 6, 1000 },
    { "THR_PWM_MAX", 7, 2000 },
    { "MSG_INR_MS",  8, 3000 },
}

assert(param:add_table(PKEY, PPFX, #PARAM_DEFS), "cannot allocate EP_ parameter table")

local P = {}
for _, def in ipairs(PARAM_DEFS) do
    assert(param:add_param(PKEY, def[2], def[1], def[3]), "cannot create parameter " .. def[1])
    local p = Parameter()
    assert(p:init(PPFX .. def[1]), "cannot bind parameter " .. PPFX .. def[1])
    P[def[1]] = p
end

-- ------------------------------------------------------- telemetry downlink

-- u16 fields are raw values scaled by 0.1 on the ECU side
local TEL_FMT = ">BHHhhhhhHHbHHBBhHHLHB"
local FRAME_LEN = 39

local EcuFrame = {}
EcuFrame.__index = EcuFrame

function EcuFrame.new()
    return setmetatable({
        ecu_status = 0,   -- ECU health flag
        thr_raw    = 0,   -- throttle position x10 (%)
        rpm        = 0,
        cht1       = 0,   -- cylinder head temp 1 (C x10)
        cht2       = 0,   -- cylinder head temp 2 (C x10)
        egt1       = 0,   -- exhaust gas temp 1 (C x10)
        egt2       = 0,   -- exhaust gas temp 2 (C x10)
        fuel_press = 0,   -- fuel pressure (mbar x10)
        inj_us     = 0,   -- injection pulse width (us)
        spark_adv  = 0,   -- ignition advance (deg, i8)
        amb_temp   = 0,
        amb_press  = 0,   -- ambient pressure (mbar)
        ecu_supply = 0,   -- ECU supply voltage (x10)
        reserved   = 0,
        ecu_5v     = 0,   -- internal 5V rail
        ecu_temp   = 0,
        flow_lh    = 0,   -- instantaneous consumption (L/h x10)
        used_raw   = 0,   -- total consumption (L x10)
        batch      = 0,   -- ECU hardware revision
        fault_bits = 0,   -- warning bitmask, bits 9..10 carry starter phase
    }, EcuFrame)
end

function EcuFrame:unpack(buf)
    if buf:len() < string.packsize(TEL_FMT) then
        return false
    end
    self.ecu_status, self.thr_raw, self.rpm,
    self.cht1, self.cht2, self.egt1, self.egt2,
    self.fuel_press, self.inj_us, self.spark_adv,
    self.amb_temp, self.amb_press,
    self.ecu_supply, self.reserved, self.ecu_5v, self.ecu_temp,
    self.flow_lh, self.used_raw, self.batch, self.fault_bits =
        string.unpack(TEL_FMT, buf)
    return true
end

-- --------------------------------------------------------- frame deframer

local ST_WAIT_H1, ST_WAIT_H2, ST_BODY, ST_CHECK = 0, 1, 2, 3

local EcuLink = {}
EcuLink.__index = EcuLink

function EcuLink.new()
    return setmetatable({
        uart    = nil,
        state   = ST_WAIT_H1,
        body    = "",
        sum     = 0,
        tx_seq  = 0,
    }, EcuLink)
end

-- feed one received byte, returns a complete payload string or nil
function EcuLink:feed(byte)
    if self.state == ST_WAIT_H1 then
        if byte == 0xAA then
            self.state = ST_WAIT_H2
        end
    elseif self.state == ST_WAIT_H2 then
        if byte == 0x55 then
            self.state = ST_BODY
            self.body = ""
            self.sum = 0
        else
            self.state = ST_WAIT_H1
        end
    elseif self.state == ST_BODY then
        self.body = self.body .. string.char(byte)
        self.sum = self.sum + byte
        if self.body:len() >= FRAME_LEN then
            self.state = ST_CHECK
        end
    elseif self.state == ST_CHECK then
        self.state = ST_WAIT_H1
        if byte == (self.sum & 0xFF) then
            return self.body
        end
    end
    return nil
end

-- pull every pending byte from the uart, return the newest decoded frame
function EcuLink:poll()
    if self.uart == nil then
        return nil
    end

    local avail = self.uart:available():toint()
    for _ = 1, avail do
        local b = self.uart:read()
        if b < 0 then
            gcs:send_text(MAV_SEVERITY.ERROR, "EP Engine: uart read failed")
            return nil
        end
        local payload = self:feed(b)
        if payload ~= nil then
            local frame = EcuFrame.new()
            if frame:unpack(payload) then
                return frame
            end
        end
    end
    return nil
end

-- build one uplink command frame (throttle percent 0..100, starter 0/1)
function EcuLink:command(throttle_pct, starter_cmd)
    local seq = self.tx_seq
    self.tx_seq = (seq + 1) & 0xFF

    local frame = string.pack(">BBHBHBB", 0x55, 0xAA,
                              math.floor(throttle_pct * 10),
                              starter_cmd, 0, 0, seq)
    local sum = 0
    for i = 3, frame:len() do
        sum = sum + frame:byte(i)
    end
    return frame, sum & 0xFF
end

-- ------------------------------------------------------------ starter logic

local PHASE_OFF, PHASE_CRANK, PHASE_RUN = 0, 1, 2

local starter_phase = PHASE_OFF
local crank_loops = 0

-- true while the starter RC channel is held above its arming threshold
local function starter_requested()
    local ch = P.STR_FUNC:get() or 0
    if ch < 1 then
        return false
    end
    local pwm = SRV_Channels:get_output_pwm(ch)
    if pwm == nil then
        return false
    end
    return pwm >= (P.STR_MIN:get() or 1800)
end

-- 1 = keep cranking this cycle, 0 = idle; gives up after 20 loops
local function starter_command(requested)
    if not arming:is_armed() or not requested or starter_phase ~= PHASE_OFF then
        crank_loops = 0
        return 0
    end

    if crank_loops == 0 then
        gcs:send_text(MAV_SEVERITY.EMERGENCY, "EP Engine: cranking")
    end

    if crank_loops >= 20 then
        crank_loops = 0
        return 0
    end

    crank_loops = crank_loops + 1
    return 1
end

-- -------------------------------------------------------------- EFI backend

local efi_backend = nil
local throttle_enabled = true

local function to_kelvin(celsius)
    return celsius + 273.15
end

-- push one decoded ECU frame into the scripting EFI backend
local function publish_telemetry(frame)
    if efi_backend == nil then
        return false
    end

    local state = EFI_State()
    local cyl = Cylinder_Status()

    state:engine_speed_rpm(uint32_t(frame.rpm))

    -- the ECU exposes two head temperatures; the second slot is mapped to
    -- the EFI exhaust gas temperature field for display purposes
    cyl:cylinder_head_temperature(to_kelvin(frame.cht1 * 0.1))
    cyl:exhaust_gas_temperature(to_kelvin(frame.cht2 * 0.1))

    state:fuel_pressure(frame.fuel_press * 0.1)             -- mbar x10 -> kPa
    state:ignition_voltage(frame.ecu_supply * 0.1)
    state:atmospheric_pressure_kpa(frame.amb_press * 0.1)   -- mbar x10 -> kPa

    -- raw accumulator holds litres x10; backend wants cm3
    state:estimated_consumed_fuel_volume_cm3(frame.used_raw * 100.0)

    -- raw flow is L/h x10; backend wants cm3/min
    state:fuel_consumption_rate_cm3pm(frame.flow_lh / 60 * 100.0)

    state:throttle_position_percent(math.floor(frame.thr_raw * 0.1))
    state:last_updated_ms(millis())

    state:cylinder_status(cyl)
    efi_backend:handle_scripting(state)
    return true
end

-- ---------------------------------------------------------- fault reporting

local last_report_ms = {}
local now_ms = 0

-- rate-limited gcs message, one slot per fault id
local function report(id, level, text)
    local interval = P.MSG_INR_MS:get() or 3000
    local last = last_report_ms[id]
    if last == nil or (now_ms - last >= interval) then
        last_report_ms[id] = now_ms
        gcs:send_text(level, text)
    end
end

local FAULT_TEXTS = {
    [0x0001] = { MAV_SEVERITY.WARNING, "EP Engine: CHT below limit" },
    [0x0002] = { MAV_SEVERITY.WARNING, "EP Engine: CHT above limit" },
    [0x0004] = { MAV_SEVERITY.WARNING, "EP Engine: CHT cylinder mismatch" },
    [0x0008] = { MAV_SEVERITY.WARNING, "EP Engine: ECU supply voltage low" },
    [0x0010] = { MAV_SEVERITY.WARNING, "EP Engine: ECU supply voltage high" },
    [0x0020] = { MAV_SEVERITY.WARNING, "EP Engine: fuel pressure low" },
    [0x0040] = { MAV_SEVERITY.WARNING, "EP Engine: fuel pressure high" },
    [0x0080] = { MAV_SEVERITY.WARNING, "EP Engine: ECU over temperature" },
    [0x0100] = { MAV_SEVERITY.WARNING, "EP Engine: ECU link lost" },
    [0x0200] = { MAV_SEVERITY.WARNING, "EP Engine: throttle actuator fault" },
}

local function report_faults(frame)
    -- bits 9..10 of the fault word mirror the starter phase reported by the ECU
    starter_phase = (frame.fault_bits >> 9) & 3

    if not frame.ecu_status then
        return
    end

    now_ms = millis():toint()

    for mask, info in pairs(FAULT_TEXTS) do
        if (frame.fault_bits & mask) ~= 0 then
            report(mask, info[1], info[2])
        end
    end
end

-- ---------------------------------------------------------------- main loop

local link = EcuLink.new()
local last_tx_ms = millis():toint()

local function discard_pending_uart()
    if link.uart == nil then
        return
    end
    local n = link.uart:available():toint()
    for _ = 1, n do
        link.uart:read()
    end
end

-- forward vehicle throttle + starter request to the ECU at 20 Hz
local function send_throttle()
    if link.uart == nil then
        return
    end

    local now = millis():toint()
    if now - last_tx_ms < 50 then
        return
    end
    last_tx_ms = now

    local pct = 0
    if throttle_enabled then
        local pwm = SRV_Channels:get_output_pwm(k_throttle_servo_fn)
        if pwm ~= nil then
            local lo = P.THR_PWM_MIN:get() or 1000
            local hi = P.THR_PWM_MAX:get() or 2000
            pwm = math.min(math.max(pwm, lo), hi)
            pct = (pwm - lo) / (hi - lo) * 100.0
        else
            gcs:send_text(MAV_SEVERITY.WARNING,
                          "EP Engine: throttle servo function has no pwm")
        end
    end

    local frame, crc = link:command(pct, starter_command(starter_requested()))
    for i = 1, frame:len() do
        if link.uart:write(frame:byte(i)) < 0 then
            return
        end
    end
    link.uart:write(crc)
end

-- throttle lock: while the lock channel sits below its threshold the ECU
-- throttle is forced to zero (engine kills instantly)
local function apply_throttle_lock()
    local ch = P.THR_FUNC:get() or 0
    if ch < 1 then
        throttle_enabled = true
        return
    end

    local pwm = SRV_Channels:get_output_pwm(ch)
    if pwm == nil then
        throttle_enabled = false
        gcs:send_text(MAV_SEVERITY.WARNING,
                      "EP Engine: throttle lock channel has no pwm")
        return
    end

    throttle_enabled = pwm > (P.THR_MIN:get() or 1500)
end

local function main_loop()
    if P.ENABLE:get() < 1 then
        return
    end

    if efi_backend == nil then
        efi_backend = efi:get_backend(0)
        if efi_backend == nil then
            gcs:send_text(MAV_SEVERITY.ERROR,
                          "EP Engine: scripting EFI backend missing (check EFI_TYPE=7)")
            return
        end
    end

    if link.uart == nil then
        local uart = serial:find_serial(0)
        if uart == nil then
            gcs:send_text(MAV_SEVERITY.ERROR,
                          "EP Engine: no free serial port for engine link")
            return
        end
        uart:begin(uint32_t(115200))
        link.uart = uart
        discard_pending_uart()
    end

    now_ms = millis():toint()

    local frame = link:poll()
    if frame ~= nil then
        publish_telemetry(frame)
        report_faults(frame)
    end

    apply_throttle_lock()
    send_throttle()
end

local function scheduler()
    local ok, err = pcall(main_loop)
    if not ok then
        gcs:send_text(MAV_SEVERITY.CRITICAL, "EP Engine: internal error: " .. err)
        return scheduler, 1000
    end
    return scheduler, 25
end

gcs:send_text(MAV_SEVERITY.INFO, "EP Engine driver script loaded")
return scheduler()