import esphome.codegen as cg
import esphome.config_validation as cv
from esphome.components import sensor, voltage_sampler
from esphome.const import (
    CONF_ID,
    CONF_GAIN,
    CONF_CHANNEL,
    DEVICE_CLASS_VOLTAGE,
    STATE_CLASS_MEASUREMENT,
    UNIT_VOLT,
)

from .. import CONF_ADS131M08_ID, ADS131M08Hub, ads131m08_ns
from ..sensor_rms import RMS_SENSOR_SCHEMA, to_code_ac

AUTO_LOAD = [
    "sensor", "voltage_sampler",
]

MAX_CHANNELS = 15
CONF_OFFSET_CALIBRATION = "offset_calibration"
CONF_GAIN_CALIBRATION = "gain_calibration"
CONF_PHASE_CALIBRATION = "phase_calibration"
CONF_INPUT_SELECT = "input_select"
CONF_DC_BLOCK = "dc_block"
ICON_CURRENT_DC = "mdi:current-dc"

DEPENDENCIES = ["ads131m08"]

GAIN = ads131m08_ns.enum("ADC_PGA_GAIN")
ALLOWED_GAINS = {
    "1": GAIN.PGA_1,
    "2": GAIN.PGA_2,
    "4": GAIN.PGA_4,
    "8": GAIN.PGA_8,
    "16": GAIN.PGA_16,
    "32": GAIN.PGA_32,
    "64": GAIN.PGA_64,
    "128": GAIN.PGA_128,
    1: GAIN.PGA_1,
    2: GAIN.PGA_2,
    4: GAIN.PGA_4,
    8: GAIN.PGA_8,
    16: GAIN.PGA_16,
    32: GAIN.PGA_32,
    64: GAIN.PGA_64,
    128: GAIN.PGA_128,
}

ADC_INPUT = ads131m08_ns.enum("ADC_INPUT_CHANNEL_MUX")
ALLOWED_MUX_INP = {
    "normal": ADC_INPUT.ICM_AIN0P_AIN0N,
    "shorted" : ADC_INPUT.ICM_INPUT_SHORTED,
    "positive_dc" : ADC_INPUT.ICM_POSITIVE_DC_TEST_SIGNAL,
    "negative_dc" : ADC_INPUT.ICM_NEGATIVE_DC_TEST_SIGNAL,
}

Channel = ads131m08_ns.class_(
    "Channel", sensor.Sensor, cg.Component, voltage_sampler.VoltageSampler
)

CONFIG_SCHEMA = (
    sensor.sensor_schema(
        Channel,
        accuracy_decimals=2,
        icon=ICON_CURRENT_DC,
        unit_of_measurement=UNIT_VOLT,
        device_class=DEVICE_CLASS_VOLTAGE,
        state_class=STATE_CLASS_MEASUREMENT,
        ).extend(
            RMS_SENSOR_SCHEMA
    )
    .extend(
        {
            cv.GenerateID(CONF_ADS131M08_ID): cv.use_id(ADS131M08Hub),
            cv.Required(CONF_CHANNEL): cv.int_range(min=0, max=MAX_CHANNELS),
            cv.Optional(CONF_GAIN, default=1): cv.enum(ALLOWED_GAINS, int=True),
            cv.Optional(CONF_OFFSET_CALIBRATION, default=0): cv.int_range(min=-8388608, max=8388607),  # should use volts, but need to figure out conversion function first
            cv.Optional(CONF_GAIN_CALIBRATION, default=1): cv.float_range(min=0, max=2),
            cv.Optional(CONF_PHASE_CALIBRATION, default=0): cv.int_range(min=-512, max=511), # should use degrees, but need to figure out conversion function first
            cv.Optional(CONF_INPUT_SELECT, default='normal'): cv.enum(ALLOWED_MUX_INP, int=False),
            cv.Optional(CONF_DC_BLOCK, default=False): cv.boolean,
        }
    )
    #.extend(cv.polling_component_schema("1s"))
    .extend(cv.COMPONENT_SCHEMA)
)

# we are using 2 sensors:
# 1. channel_sensor: this represents 1 of the 8 ads131m08 channels and is used to program the adc channel gain, offset calibration, etc.
#    This sensor publishes instantaneous sampled value and can be used for DC measurements
# 2. ac_sensor: to publish rms ac value; only use if absolutely necessary as it it expensive in terms of loop time
async def to_code(config):
    channel_sensor = cg.new_Pvariable(config[CONF_ID])
    await cg.register_parented(channel_sensor, config[CONF_ADS131M08_ID])
    channel = config[CONF_CHANNEL]
    cg.add(channel_sensor.set_channel_number(channel))
    gain = config[CONF_GAIN]
    cg.add(channel_sensor.set_gain(gain))
    mux_inp = config[CONF_INPUT_SELECT]
    cg.add(channel_sensor.set_mux_input(mux_inp))
    gain_cal = config[CONF_GAIN_CALIBRATION]
    cg.add(channel_sensor.set_gain_calibration(gain_cal))
    offset_cal = config[CONF_OFFSET_CALIBRATION]
    cg.add(channel_sensor.set_offset_calibration(offset_cal))
    phase_cal = config[CONF_PHASE_CALIBRATION]
    cg.add(channel_sensor.set_phase_calibration(phase_cal))
    dc_block = config[CONF_DC_BLOCK]
    cg.add(channel_sensor.set_dc_block(dc_block))
    hub = await cg.get_variable(config[CONF_ADS131M08_ID])
    cg.add(hub.register_sensor_dc(channel, channel_sensor))
    await sensor.register_sensor(channel_sensor, config)
    await cg.register_component(channel_sensor, config)
    if ac_sensor := await to_code_ac(config):
        await cg.register_parented(ac_sensor, channel_sensor)