#!/usr/bin/env python3
import asyncio
import json
import os
import sys
import time
from utils import ChangedFloat
from ws.message import MessageDispatcher
from ws.server.ws_server_multi_user import WsServerMultiUser
from components.sensor import TempSensorFactory
from components.pid import PidFactory
from components.plant import Pot
from components.actor import HeaterFactory, StirrerFactory
from components.sud import Sud
from components.sud_forecast import SudForecastEstimator
from tasks import TaskManager, TempSensorTask, HeaterTask, PotTask, TcTask, StirrerTask, SudTask, SudLogTask

import argparse as ap


class Tee:
	"""Duplicates writes to multiple streams - used to mirror stdout/stderr
	(everything the rest of the codebase already reaches via print()) into
	a log file under logs/ as well as the console, without having to touch
	every print() call site."""

	def __init__(self, *streams):
		self.streams = streams

	def write(self, data):
		for stream in self.streams:
			stream.write(data)

	def flush(self):
		for stream in self.streams:
			stream.flush()


if __name__ == '__main__':
	parser = ap.ArgumentParser()
	parser.add_argument("-c", "--config", default="config.json")
	parser.add_argument("-m", "--logdir", default="logs")

	args = parser.parse_args()

	os.makedirs(args.logdir, exist_ok=True)
	log_path = f"{args.logdir}/brewpi.{time.strftime("%Y%m%d%H%M%S", time.localtime())}.log"
	log_file = open(log_path, "a", buffering=1)
	sys.stdout = Tee(sys.stdout, log_file)
	sys.stderr = Tee(sys.stderr, log_file)
	print("Logging to {}".format(log_path))

	config = json.load(open(args.config))
	dispatcher = MessageDispatcher()
	server = WsServerMultiUser(listener=dispatcher)
	taskmgr = TaskManager()

	DT = config['Controller']['dt']
	DT_TASK = 1.0 / config['Controller']['sim_warp_factor']
	theta_amb = config['ambient_temperature']
	plant_sim = "sim" in config['Controller']['plant_name']

	# Mash schedule - starts out empty; a client loads one of sude/*.json's
	# several schedules onto it later (see components/sud.py's Sud).
	sud = Sud()

	# Sensor
	sensor = TempSensorFactory.create(config['Controller']['sensor_name'], temp_offset=-0.15, variance=0.01)
	sensor_task = TempSensorTask(sensor, DT_TASK, dispatcher.msgio_get("Sensor"))
	taskmgr.add(sensor_task)

	# Plant - deliberately left without plant params (M/C/L/Td) here: a
	# Sud's own doc is the only source for those now (see
	# tasks/sud.py's SudTask.apply_plant_params(), called immediately on
	# Load), so the plant stays inert (Pot.is_configured() is False,
	# PotTask skips process() - see there) until one actually is. Ambient
	# is independent of any Sud (a global setting, changeable live via
	# the System channel - see on_system_recv() below), so it's set
	# right away regardless.
	pot = Pot(DT)
	pot.set_ambient_temperature(theta_amb)
	taskmgr.add(PotTask(pot, DT_TASK, dispatcher.msgio_get("Pot")))

	# Heater
	heater = HeaterFactory.create(config['Controller']['heater_name'], config['Heater'])
	heater.set_on_changed("power_eff", ChangedFloat(pot.set_power, prec=0).set)
	heater_task = HeaterTask(heater, DT_TASK, dispatcher.msgio_get("Heater"))
	taskmgr.add(heater_task)

	# Temperature Controller - PID gains are server/hardware-level
	# config, independent of any Sud, so they're set right away; "Normal"
	# has no internal model/ambient at all (see temp_controller.py vs.
	# temp_controller_smith.py.), and Smith's model plant params are -
	# like the real Pot's above - deliberately left unset here, only
	# ever coming from a Sud's own doc.
	tc = PidFactory.create(config['Controller']['pid_type'], DT)
	tc.set_params(config['TempCtrl'])
	if hasattr(tc, 'set_ambient_temperature'):
		tc.set_ambient_temperature(theta_amb)
	tc_task = TcTask(tc, DT_TASK, dispatcher.msgio_get("TempCtrl"))
	taskmgr.add(tc_task)

	# Stirrer
	stirrer = StirrerFactory.create(config['Controller']['stirrer_name'], DT, config['Stirrer'])
	stirrer_task = StirrerTask(stirrer, DT_TASK, dispatcher.msgio_get("Stirrer"))
	taskmgr.add(stirrer_task)

	# Predicts how long a loaded schedule will actually take, by simulating
	# it with the same kind of plant/controller (and params) as above -
	# see components/sud_forecast.py.
	forecast_estimator = SudForecastEstimator(
		DT, theta_amb, config['Controller']['pid_type'], config['TempCtrl'], heater.get_power_max())
	sud_task = SudTask(sud, tc, stirrer, pot, DT, DT_TASK, dispatcher.msgio_get("Sud"), forecast_estimator)
	taskmgr.add(sud_task)
	# Records every Sud run's measured data and forecast to their own
	# JSON files under logs/, for later offline analysis - see
	# tasks/sud_log.py.
	taskmgr.add(SudLogTask(sud, tc, heater, heater_task, sud_task, DT_TASK))

	# Assign data flow
	# Assign tc control value to heater
	tc.set_on_changed("y", heater_task.actor)

	# Assign temp. sensor readings to tc
	sensor_task.set_on_changed("temp", ChangedFloat(tc.set_theta_ist, prec=2).set)

	# Assign heater power set to tc model - "Normal" has no internal model
	if hasattr(tc, "set_model_power"):
		heater.set_on_changed("power_set", tc.set_model_power)

	# For simulation
	if "sim" in config['Controller']['sensor_name']:
		pot.set_on_changed("temp", ChangedFloat(sensor.set_fake_temp, prec=3).set)
		sensor.set_fake_temp(pot.temp)

	# Static, global info that doesn't belong to any one task - sent once;
	# a client connecting later still gets it via the dispatcher's replay
	# of accumulated global_state on subscribe. AmbientTemp can also be
	# changed live by a client (e.g. the GUI's ambient temperature field).
	msg_system = dispatcher.msgio_get("System")

	async def send_system_info():
		await msg_system.send({'AmbientTemp': theta_amb, 'WarpFactor': DT / DT_TASK, 'PlantSim': plant_sim})
	asyncio.ensure_future(send_system_info())

	async def on_system_recv(data):
		global theta_amb
		if 'AmbientTemp' in data:
			theta_amb = data['AmbientTemp']
			pot.set_ambient_temperature(theta_amb)
			if hasattr(tc, 'set_ambient_temperature'):
				tc.set_ambient_temperature(theta_amb)
			forecast_estimator.set_ambient_temperature(theta_amb)
			await msg_system.send({'AmbientTemp': theta_amb})
	msg_system.set_recv_handler(on_system_recv)

	# Message dispatcher
	h_dispatcher = taskmgr.start()
	h_server = server.listen("0.0.0.0", 8765)
	asyncio.gather(h_dispatcher, h_server)
	server.run_forever()