#!/usr/bin/env python3
import asyncio
import functools
import json
import os
import sys
import threading
import time
from http.server import SimpleHTTPRequestHandler, ThreadingHTTPServer
from utils import ChangedFloat
from ws.message import MessageDispatcher
from ws.server.ws_server_multi_user import WsServerMultiUser
from components.pid import PidFactory
from components.plant import PlantFactory
from components.actor import 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")
	# Simulated seconds per tick, and how many of those fit into one real
	# second - run-mode knobs (how fast to drive this particular process),
	# not config.json material (a *plant/hardware* description that stays
	# the same regardless of how this run happens to be invoked) - see
	# tasks/sud_forecast.py's SudForecastEstimator and every *Task's own
	# dt/interval split for where these actually end up. Defaulting both
	# to 1.0 means real time, 1-simulated-second ticks unless asked
	# otherwise - e.g. a much higher --sim-warp-factor for dev/testing.
	parser.add_argument("--dt", type=float, default=1.0)
	parser.add_argument("--sim-warp-factor", type=float, default=1.0)
	# Serves web/ (the browser client - see web/README or the project
	# README's "Browser client" section) so a plain browser can reach this
	# same rig remotely, no desktop GUI required. A stdlib HTTP server in
	# its own thread, deliberately independent of the asyncio WebSocket
	# server below (ws/server/ws_server.py) - neither can affect the
	# other's behavior or availability.
	parser.add_argument("--http-port", type=int, default=8080)

	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 = args.dt
	DT_TASK = 1.0 / args.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()

	# Heater/Pot/Sensor - built together as one consistent rig by a single
	# Controller.plant_name, rather than three independently-named
	# components that in practice are never actually mixed and matched
	# (see components/plant/plant_factory.py): "sim" gets HeaterSim/Pot
	# (the modeled plant)/TempSensorSim; anything else gets HeaterHendi/
	# PotReal (no model - the real kettle does its own physics)/
	# TempSensor_max31865.
	heater, pot, sensor = PlantFactory.create(config['Controller']['plant_name'], DT, config['Heater'])

	sensor_task = TempSensorTask(sensor, DT_TASK, dispatcher.msgio_get("Sensor"))
	taskmgr.add(sensor_task)

	# Pot - 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
	# simulated Pot stays inert (Pot.is_configured() is False, PotTask
	# skips process() - see there) until one actually is loaded (PotReal
	# just ignores this - see there). 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.set_ambient_temperature(theta_amb)
	taskmgr.add(PotTask(pot, DT_TASK, dispatcher.msgio_get("Pot")))

	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'])
	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_powers(), args.sim_warp_factor)
	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)
	# HeaterTask owns TC enable: enabled in closed-loop mode, disabled in open-loop.
	heater_task.set_on_closed_loop_changed(tc.set_enabled)
	# Brew end (DONE/IDLE): shut heater off and broadcast open-loop/power-0 to clients.
	sud_task.set_on_end(heater_task.shutdown)

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

	heater.set_on_changed("power_set", tc.set_model_power)

	# For simulation - the sim sensor has no real plant to read, so it
	# just reports back whatever the modeled Pot's own temperature is.
	if plant_sim:
		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)
			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)

	web_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), os.pardir, "web")
	http_handler = functools.partial(SimpleHTTPRequestHandler, directory=web_dir)
	http_server = ThreadingHTTPServer(("0.0.0.0", args.http_port), http_handler)
	threading.Thread(target=http_server.serve_forever, daemon=True).start()
	print("Serving {} on http://0.0.0.0:{}".format(web_dir, args.http_port))

	loop = asyncio.get_event_loop()
	try:
		loop.run_forever()
	except KeyboardInterrupt:
		print("\nShutting down...")
	finally:
		# Cancel all tasks so their finally blocks (e.g. HeaterTask's
		# `with device.open()`) get a chance to run before we close the port.
		pending = asyncio.all_tasks(loop)
		for task in pending:
			task.cancel()
		if pending:
			loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True))
		# Belt-and-suspenders: explicitly close the heater connection even if
		# the task cleanup above failed to reach HeaterHendi.activate(False).
		heater.close()
		http_server.shutdown()
		loop.close()
		print("Server stopped.")