/* Copyright 2023-2025 Bruce Merry
 *
 * This program is free software: you can redistribute it and/or modify it
 * under the terms of the GNU General Public License as published by the Free
 * Software Foundation, either version 3 of the License, or (at your option)
 * any later version.
 *
 * This program is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
 * FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for
 * more details.
 *
 * You should have received a copy of the GNU General Public License along
 * with this program. If not, see <https://www.gnu.org/licenses/>.
 */

use async_trait::async_trait;
use chrono::{DateTime, Duration, Utc};
use futures::StreamExt;
use log::{error, info, warn};
use std::cmp::min;
use std::collections::VecDeque;
use std::sync::Mutex;
use std::time::Instant;
use tokio::time::MissedTickBehavior;
use tokio_stream::StreamMap;
use tokio_util::sync::CancellationToken;

use crate::config::{CoilConfig, Config, InverterConfig};
use crate::esp_api::{API, ScheduleResponse};
use crate::inverter::{Info, Inverter, Result};
use crate::monitoring::{CoilUpdate, Monitor, SocUpdate};
use crate::sun::SolarPredictor;

pub struct State {
    pub response: ScheduleResponse,
    pub time: DateTime<Utc>,
}

pub async fn poll_esp(
    api: &API,
    schedule_id: &str,
    interval: std::time::Duration,
    state: &Mutex<Option<State>>,
    token: CancellationToken,
) {
    let mut interval = tokio::time::interval(interval);
    interval.set_missed_tick_behavior(MissedTickBehavior::Delay);
    loop {
        tokio::select! {
            _ = interval.tick() => {},
            _ = token.cancelled() => { break; }
        }
        match api.schedule(schedule_id).await {
            Ok(response) => {
                let mut lock = state.lock().unwrap();
                *lock = Some(State {
                    response,
                    time: Utc::now(),
                });
                drop(lock);
                info!("Successfully updated schedule info from EskomSePush");
            }
            Err(err) => {
                warn!("Failed to update from EskomSePush: {err:?}");
            }
        }
    }
}

fn filter_state(state: &Option<State>, min_time: DateTime<Utc>) -> Option<&State> {
    state.as_ref().filter(|state| state.time >= min_time)
}

/// Number of (non-integer) hours in a duration
fn duration_hours(duration: Duration) -> f64 {
    (duration.num_milliseconds() as f64) / 3600000.0
}

/// What to simulate when no load-shedding and not enough solar
enum SimMode {
    /// Power drains from battery
    Drain,
    /// Battery level held steady
    Hold,
    /// Charge battery as fast as possible
    Charge,
}

/// What to simulate when modelling export
enum ExportMode {
    /// Export first, including from the battery
    Drain,
    /// Export first, but do not export from battery
    Hold,
}

fn target_soc_helper(
    config: &InverterConfig,
    predictor: &dyn SolarPredictor<Utc>,
    state: &State,
    info: &Info,
    now: DateTime<Utc>,
    mode: SimMode,
) -> (f64, DateTime<Utc>) {
    let step = Duration::seconds(60);
    let step_h = duration_hours(step);
    let depth = info.capacity - config.min_soc * 0.01 * info.capacity;

    let mut base_wh = 0.0;
    let mut worst = 0.0_f64;
    let mut floor = -depth;
    let mut worst_time = now;
    /* Project battery level forward for 24 hours, using optimistic
     * assumptions about solar PV and consumption. Whenever the
     * current point falls into load-shedding, check that there will
     * be enough to get to the end with pessimistic assumptions.
     */
    let goal = now + Duration::seconds(86400);
    let mut t = now;
    let mut observe = |wh, t| {
        if wh < worst {
            worst = wh;
            worst_time = t;
        }
    };
    while t < goal {
        let mut have_grid = true;
        for event in state.response.events.iter() {
            if t >= event.start && t < event.end {
                have_grid = false;
                let end_wh = base_wh - config.max_discharge_power * duration_hours(event.end - t);
                observe(end_wh.max(floor), t);
            }
        }
        let mut power = predictor.power(&(t + step / 2));
        power -= config.min_discharge_power;
        power = power.min(config.charge_power);
        if have_grid {
            power = match mode {
                SimMode::Drain => power,
                SimMode::Hold => power.max(0.0),
                SimMode::Charge => config.charge_power,
            };
        }
        base_wh += power * step_h;
        t += step;

        floor = floor.max(base_wh - depth);
        observe(base_wh.max(floor), t);
    }

    let extra = -worst / info.capacity * 100.0;
    let target = config.min_soc + extra;
    let target = target.clamp(0.0, 100.0);
    (target, worst_time)
}

fn target_socs(
    config: &InverterConfig,
    predictor: &dyn SolarPredictor<Utc>,
    state: Option<&State>,
    info: &Info,
    now: DateTime<Utc>,
) -> (f64, f64, f64) {
    match state {
        None => (config.fallback_soc, config.fallback_soc, config.min_soc),
        Some(state) => {
            for event in state.response.events.iter() {
                info!("Load-shedding from {} to {}", event.start, event.end);
            }
            let (target_high, _) =
                target_soc_helper(config, predictor, state, info, now, SimMode::Drain);
            let (target_low, _) =
                target_soc_helper(config, predictor, state, info, now, SimMode::Hold);
            let (alarm, _) =
                target_soc_helper(config, predictor, state, info, now, SimMode::Charge);
            (target_low, target_high, alarm)
        }
    }
}

fn export_soc_helper(
    config: &InverterConfig,
    predictor: &dyn SolarPredictor<Utc>,
    info: &Info,
    now: DateTime<Utc>,
    mode: ExportMode,
) -> f64 {
    // TODO: take load shedding into account (can't export during load shedding)
    let step = Duration::seconds(60);
    let step_h = duration_hours(step);
    let goal = now + Duration::seconds(86400);
    let mut t = now;
    let mut cur = 0.0f64;
    let mut peak = 0.0f64;
    let mut entered = false;

    while t < goal {
        let power = predictor.power(&(t + step / 2));
        let generating = power > 0.0;
        let power = power - config.min_discharge_power;
        let power = match mode {
            ExportMode::Drain => power - info.export_power,
            ExportMode::Hold => {
                if power < 0.0 {
                    power // Not generating enough - battery runs down
                } else {
                    (power - info.export_power).max(0.0) // Export, and charge battery with leftover
                }
            }
        };
        // TODO: also need to clamp rate at which battery can deliver power.
        // For now assume the battery can supply the maximum export power.
        let power = power.min(config.charge_power);
        cur += power * step_h;
        peak = peak.max(cur);
        t += step;
        if generating {
            entered = true;
        } else if entered {
            break; // Reached the end of production for the day
        }
    }

    (1.0 - peak / info.capacity) * 100.0
}

fn export_soc(
    config: &InverterConfig,
    predictor: &dyn SolarPredictor<Utc>,
    info: &Info,
    now: DateTime<Utc>,
) -> (f64, f64) {
    (
        export_soc_helper(config, predictor, info, now, ExportMode::Hold),
        export_soc_helper(config, predictor, info, now, ExportMode::Drain),
    )
}

async fn update_soc(
    inverter: &mut dyn Inverter,
    config: &InverterConfig,
    predictor: &dyn SolarPredictor<Utc>,
    monitor: &mut dyn Monitor,
    state: &Mutex<Option<State>>,
    esp_timeout: Duration,
    now: DateTime<Utc>,
) -> Result<()> {
    let info = inverter.get_info().await?;
    let current_soc = inverter.get_soc().await?;
    let net_production = inverter.get_net_production().await?;
    let mut target;
    let update;
    let full_export;

    {
        let guard = &state.lock().unwrap();
        let state = filter_state(guard, now - esp_timeout);
        let est_start = Instant::now();
        let (target_soc_low, target_soc_high, alarm_soc) =
            target_socs(config, predictor, state, &info, now);
        let (target_soc_export_low, target_soc_export_high) =
            export_soc(config, predictor, &info, now);
        info!(
            "Target SoC range is {:.2} - {:.2} (alarm at {:.2}); export range is {:.2} - {:.2}; computed in {:.3} s",
            target_soc_low,
            target_soc_high,
            alarm_soc,
            target_soc_export_low,
            target_soc_export_high,
            est_start.elapsed().as_secs_f64()
        );
        target = current_soc.min(target_soc_high).max(target_soc_low);
        full_export = if info.export_enabled && config.control_export {
            let low = target_soc_export_low.max(target_soc_low);
            let high = target_soc_export_high.max(target_soc_high);
            // Turn off full-export when battery is full, because Sunsynk has
            // a bug where non-essential loads count against export limit when
            // full-export is enabled.
            if current_soc < 99.0
                && (current_soc >= high || (current_soc >= low && net_production > 0.0))
            {
                target = target.max(current_soc.min(high));
                Some(true)
            } else {
                Some(false)
            }
        } else {
            None
        };

        let mut is_loadshedding = false;
        let mut next_change = None;
        if let Some(state) = state {
            for event in state.response.events.iter() {
                if now >= event.start && now < event.end {
                    is_loadshedding = true;
                    next_change = Some(event.end);
                    break;
                } else if now < event.start {
                    next_change = Some(next_change.map_or(event.start, |t| min(t, event.start)));
                }
            }
        }

        update = SocUpdate {
            time: now,
            target_soc_low,
            target_soc_high,
            alarm_soc,
            target_soc_export_low,
            target_soc_export_high,
            current_soc,
            predicted_pv: predictor.power(&now),
            is_loadshedding,
            next_change,
        };
    }

    info!("Setting min SoC to {target:.2}");
    inverter.set_min_soc(target, config.fallback_soc).await?;
    if let Some(full_export) = full_export {
        info!("Setting full export to {full_export}");
        inverter.set_full_export(full_export).await?;
    }
    if let Err(err) = monitor.soc_update(update).await {
        warn!("Failed to update monitoring: {err}");
    }

    Ok(())
}

#[async_trait]
trait Controller: Send + Unpin {
    fn interval(&self) -> std::time::Duration;
    async fn update(
        &mut self,
        inverter: &mut dyn Inverter,
        monitor: &mut dyn Monitor,
        now: DateTime<Utc>,
    );
    async fn shutdown(&mut self, inverter: &mut dyn Inverter);
}

struct SocController<'a> {
    config: &'a InverterConfig,
    predictor: &'a dyn SolarPredictor<Utc>,
    state: &'a Mutex<Option<State>>,
    esp_timeout: Duration,
}

impl<'a> SocController<'a> {
    fn new(
        config: &'a InverterConfig,
        predictor: &'a dyn SolarPredictor<Utc>,
        state: &'a Mutex<Option<State>>,
        esp_timeout: Duration,
    ) -> Self {
        Self {
            config,
            predictor,
            state,
            esp_timeout,
        }
    }
}

#[async_trait]
impl Controller for SocController<'_> {
    fn interval(&self) -> std::time::Duration {
        std::time::Duration::from_secs(60)
    }

    async fn update(
        &mut self,
        inverter: &mut dyn Inverter,
        monitor: &mut dyn Monitor,
        now: DateTime<Utc>,
    ) {
        if let Err(err) = update_soc(
            inverter,
            self.config,
            self.predictor,
            monitor,
            self.state,
            self.esp_timeout,
            now,
        )
        .await
        {
            warn!("Failed to update inverter: {err}");
        }
    }

    async fn shutdown(&mut self, inverter: &mut dyn Inverter) {
        info!(
            "Shutting down, setting minimum SoC to {}",
            self.config.fallback_soc
        );
        match inverter
            .set_min_soc(self.config.fallback_soc, self.config.fallback_soc)
            .await
        {
            Ok(_) => {}
            Err(err) => {
                error!("Failed to set minimum SoC: {err}");
            }
        }
    }
}

struct CoilController<'a> {
    history: VecDeque<Option<f64>>,
    config: &'a CoilConfig,
    last_setting: Option<f64>,
}

impl<'a> CoilController<'a> {
    const CAPACITY: usize = 11;

    fn new(config: &'a CoilConfig) -> Self {
        Self {
            history: VecDeque::with_capacity(Self::CAPACITY),
            config,
            last_setting: None,
        }
    }

    async fn update_fallible(
        &mut self,
        inverter: &mut dyn Inverter,
        monitor: &mut dyn Monitor,
        now: DateTime<Utc>,
    ) -> Result<()> {
        let info = inverter.get_coil().await?;
        let mut target = None;
        if let Some(value) = &info {
            let ne = value.coil - value.inverter;
            if ne <= self.config.power_threshold {
                // It's fake power from misreading coil
                target = Some(ne + self.config.trickle);
            }
        }
        if self.history.len() == Self::CAPACITY {
            self.history.pop_front();
        }
        self.history.push_back(target);
        if self.history.len() != Self::CAPACITY {
            return Ok(());
        }
        // Compute the sum if all elements are not None
        let sum = self.history.iter().cloned().sum::<Option<f64>>();
        let mean = sum.map(|x| x / (self.history.len() as f64));
        let coil_active = info.is_some_and(|x| x.coil_active);
        if coil_active {
            if let Some(target) = mean {
                if self.last_setting.is_none_or(|x| (x - target).abs() >= 10.0) {
                    info!("Setting trickle to {target}.");
                    inverter.set_trickle(target).await?;
                    self.last_setting = Some(target);
                } else {
                    info!("Ideal trickle setting is {target}, but not setting due to hysteresis.");
                }
            } else {
                info!("Not adjusting trickle because there is no target information.")
            }
        } else {
            info!("Not adjusting trickle because coil is not active.");
        }
        let update = CoilUpdate {
            time: now,
            active: coil_active,
            target: mean,
            setting: self.last_setting,
        };
        if let Err(err) = monitor.coil_update(update).await {
            warn!("Failed to update monitoring: {err}");
        }
        Ok(())
    }
}

#[async_trait]
impl Controller for CoilController<'_> {
    fn interval(&self) -> std::time::Duration {
        std::time::Duration::from_secs(10)
    }

    async fn update(
        &mut self,
        inverter: &mut dyn Inverter,
        monitor: &mut dyn Monitor,
        now: DateTime<Utc>,
    ) {
        match self.update_fallible(inverter, monitor, now).await {
            Ok(_) => {}
            Err(err) => {
                error!("Failed to update CT coil: {err}");
            }
        }
    }

    async fn shutdown(&mut self, _inverter: &mut dyn Inverter) {}
}

pub async fn control_inverter(
    inverter: &mut dyn Inverter,
    config: &Config,
    predictor: &dyn SolarPredictor<Utc>,
    monitor: &mut dyn Monitor,
    state: &Mutex<Option<State>>,
    esp_timeout: Duration,
    token: CancellationToken,
) {
    let mut controllers: Vec<Box<dyn Controller>> = Vec::new();
    controllers.push(Box::new(SocController::new(
        &config.inverter,
        predictor,
        state,
        esp_timeout,
    )));
    if let Some(coil_config) = &config.coil {
        controllers.push(Box::new(CoilController::new(coil_config)));
    }
    let mut stream = StreamMap::new();
    for (i, controller) in controllers.iter().enumerate() {
        let mut interval = tokio::time::interval(controller.interval());
        interval.set_missed_tick_behavior(MissedTickBehavior::Delay);
        stream.insert(i, tokio_stream::wrappers::IntervalStream::new(interval));
    }

    loop {
        tokio::select! {
            Some((idx, _)) = stream.next() => { controllers[idx].update(inverter, monitor, Utc::now()).await; }
            _ = token.cancelled() => { break; }
        }
    }

    for controller in controllers.iter_mut() {
        controller.shutdown(inverter).await;
    }
}

#[cfg(test)]
mod test {
    use approx::assert_relative_eq;
    use chrono::{DateTime, Duration, NaiveDate, NaiveTime, Utc};
    use std::sync::Mutex;

    use super::{Controller, SocController, State};
    use crate::config::InverterConfig;
    use crate::esp_api::ScheduleResponse;
    use crate::inverter::test::TestInverter;
    use crate::monitoring::NullMonitor;
    use crate::sun::SolarPredictor;
    use crate::sun::test::LinearSolarPredictor;

    /// Inverter fixture, with no panels
    fn inverter_config() -> InverterConfig {
        InverterConfig {
            device: "dummy".to_string(),
            id: 1,
            min_soc: 25.0,
            fallback_soc: 50.0,
            min_discharge_power: 100.0,  // 2% per hour
            max_discharge_power: 1000.0, // 20% per hour
            charge_power: 2500.0,        // 50% per hour
            control_export: true,
            dry_run: false,
            panels: Vec::new(),
        }
    }

    /// Fixture with no loadshedding
    fn response_no_loadshedding() -> ScheduleResponse {
        use crate::esp_api::Schedule;
        ScheduleResponse {
            events: Vec::new(),
            name: "test schedule".to_string(),
            schedule: Schedule {
                days: Vec::new(),
                schedule_name: "test".to_string(),
                source: "data".to_string(),
            },
        }
    }

    /// Fixture with a single loadshedding window
    fn response_loadshedding1(
        start_hour: u32,
        start_min: u32,
        start_sec: u32,
        end_hour: u32,
        end_min: u32,
        end_sec: u32,
    ) -> ScheduleResponse {
        let mut response = response_no_loadshedding();
        response.events.push(crate::esp_api::Event {
            start: hms(start_hour, start_min, start_sec),
            end: hms(end_hour, end_min, end_sec),
            note: "".to_string(),
        });
        response
    }

    fn hms(hour: u32, min: u32, sec: u32) -> DateTime<Utc> {
        NaiveDate::from_ymd_opt(2025, 10, 11)
            .unwrap()
            .and_hms_opt(hour, min, sec)
            .unwrap()
            .and_utc()
    }

    fn predictor() -> impl SolarPredictor<Utc> {
        LinearSolarPredictor::new(
            NaiveTime::from_hms_opt(9, 0, 0).unwrap(),
            NaiveTime::from_hms_opt(15, 0, 0).unwrap(),
            3000.0,
            3000.0,
        )
    }

    /// Test that when no load-shedding data is available, the fallback SoC is used
    #[tokio::test]
    async fn test_soc_controller_fallback() {
        let mut inverter = TestInverter::new();
        let config = inverter_config();
        let state_mutex = Mutex::new(None);
        let predictor = predictor();
        let mut controller =
            SocController::new(&config, &predictor, &state_mutex, Duration::seconds(3600));
        let mut monitor = NullMonitor {};
        controller
            .update(&mut inverter, &mut monitor, hms(6, 0, 0))
            .await;
        assert_eq!(inverter.target_soc, 50.0);
    }

    /// Test that min SoC ramps down towards the start of sunlight
    #[tokio::test]
    async fn test_soc_controller_rampdown() {
        let mut inverter = TestInverter::new();
        let config = inverter_config();
        let now = hms(6, 0, 0);
        let state_mutex = Mutex::new(Some(State {
            response: response_no_loadshedding(),
            time: now,
        }));
        let predictor = predictor();
        let mut controller =
            SocController::new(&config, &predictor, &state_mutex, Duration::seconds(3600));
        let mut monitor = NullMonitor {};
        // 3 hours until the solar switches on, at 100W (2% per hour)
        controller.update(&mut inverter, &mut monitor, now).await;
        assert_eq!(inverter.target_soc, 31.0);
        // Ensure it doesn't charge up to this level if already below
        inverter.soc = 28.0;
        controller.update(&mut inverter, &mut monitor, now).await;
        assert_eq!(inverter.target_soc, 28.0);
    }

    /// Test that load-shedding ramps up the min SoC prior to load-shedding
    #[tokio::test]
    async fn test_soc_loadshedding() {
        let mut inverter = TestInverter::new();
        let config = inverter_config();
        let state_mutex = Mutex::new(Some(State {
            response: response_loadshedding1(10, 0, 0, 12, 0, 0),
            time: hms(0, 0, 0),
        }));
        let predictor = predictor();
        let mut controller =
            SocController::new(&config, &predictor, &state_mutex, Duration::seconds(86400));
        let mut monitor = NullMonitor {};
        // We could use up to 40% during loadshedding (20%/h for 2h), and want
        // to end with 25%, so we need to start load-shedding at 65%.
        controller
            .update(&mut inverter, &mut monitor, hms(10, 0, 0))
            .await;
        assert_eq!(inverter.target_soc, 65.0);
        // We charge at 50% per hour, so 0.5h earlier we must be at 40%.
        controller
            .update(&mut inverter, &mut monitor, hms(9, 30, 0))
            .await;
        assert_eq!(inverter.target_soc, 40.0);
        // As load shedding progresses, we need to preserve less
        controller
            .update(&mut inverter, &mut monitor, hms(11, 0, 0))
            .await;
        assert_eq!(inverter.target_soc, 45.0);
    }

    /// Test that we sell remaining battery if needed to make room for charging
    #[tokio::test]
    async fn test_sell_battery() {
        let mut inverter = TestInverter::new();
        // Solar model is 3kW output, so this leaves us with 500W excess (after
        // the 100W minimum load) we want to store, or 10%/h, over 6h, for 60%.
        inverter.info.export_power = 2400.0;
        let config = inverter_config();
        let state_mutex = Mutex::new(Some(State {
            response: response_no_loadshedding(),
            time: hms(0, 0, 0),
        }));
        let predictor = predictor();
        let mut controller =
            SocController::new(&config, &predictor, &state_mutex, Duration::seconds(86400));
        let mut monitor = NullMonitor {};

        controller
            .update(&mut inverter, &mut monitor, hms(9, 0, 0))
            .await;
        assert_eq!(inverter.full_export, true);
        assert_relative_eq!(inverter.target_soc, 40.0, epsilon = 1e-9);

        // We can only guarantee draining the battery at 50%/h
        controller
            .update(&mut inverter, &mut monitor, hms(8, 30, 0))
            .await;
        assert_eq!(inverter.full_export, true);
        assert_relative_eq!(inverter.target_soc, 65.0, epsilon = 1e-9);

        // If the battery is already below the target, do not try to export
        controller
            .update(&mut inverter, &mut monitor, hms(8, 0, 0))
            .await;
        assert_eq!(inverter.full_export, false);
        assert_relative_eq!(inverter.target_soc, 27.0, epsilon = 1e-9);
    }
}
