#!/usr/bin/env python3
import asyncio
import functools
import json
import os
import signal
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, ServerLogTask, 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.

	Also stamps every line with a "<date>T<time>:" prefix (existing print()
	call sites already start their message with "Component: ..." - see e.g.
	tasks/stirrer.py - so the combined line reads
	"<date>T<time>:<component>: <message>"), so log entries can be
	correlated with the JSON sample logs (ServerLogTask/SudLogTask) without
	relying on journalctl's own timestamps, which aren't present in the
	mirrored log files themselves."""

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

	def write(self, data):
		if not data:
			return
		out = []
		for line in data.splitlines(keepends=True):
			if self._at_line_start:
				out.append(time.strftime("%Y-%m-%dT%H:%M:%S:", time.localtime()))
			out.append(line)
			self._at_line_start = line.endswith('\n')
		text = ''.join(out)
		for stream in self.streams:
			stream.write(text)

	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)
	# Fixed (no-date) name written alongside the dated one - always
	# truncated and re-mirrored, so a tail-style consumer can point at one
	# unchanging path instead of tracking this run's own timestamped
	# filename (same principle as ServerLogTask/SudLogTask's log_latest*.json).
	latest_log_path = f"{args.logdir}/brewpi.latest.log"
	latest_log_file = open(latest_log_path, "w", buffering=1)
	sys.stdout = Tee(sys.stdout, log_file, latest_log_file)
	sys.stderr = Tee(sys.stderr, log_file, latest_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
	theta_amb = config['ambient_temperature']
	plant_sim = config['Heater'].get('type', 'sim') == 'sim'
	# sim_warp_factor only reciprocally scales *wait* time between task
	# ticks (DT_TASK) - it has no meaning against real hardware, which
	# ticks at its own physical pace regardless of how this flag is set,
	# so real runs are always paced at DT_TASK == 1.0 (true real time).
	sim_warp_factor = args.sim_warp_factor if plant_sim else 1.0
	DT_TASK = 1.0 / sim_warp_factor

	# 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(config.get('Pot', {}))

	# Heater/Pot/Sensor - built together as one consistent rig; each
	# component's type comes from its own config section. Heater.type==sim
	# gets HeaterSim/Pot (the modelled plant)/TempSensorSim; anything else
	# (e.g. hendi) gets HeaterHendi/PotReal/whatever TempSensor.type names.
	heater, pot, sensor = PlantFactory.create(DT, config['Heater'], config.get('TempSensor', {}))

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

	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)

	tc = PidFactory.create(config['TempCtrl']['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)

	# Seed plant params from config so pot/tc are configured before any sud
	# is loaded - Pot.water_mass in config is how much water is already in
	# the kettle at startup (e.g. strike water pre-filled). A loaded sud's
	# own apply_plant_params() will override this on the first step.
	startup_water = config.get('Pot', {}).get('water_mass', 0)
	startup_params = None
	if startup_water > 0:
		startup_params = sud.derive_plant_params(0, startup_water)
		pot.set_plant_params(startup_params)
		tc.set_model_plant_params(startup_params)

	# Stirrer
	stirrer = StirrerFactory.create(config['Stirrer']['type'], 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['TempCtrl']['pid_type'], config['TempCtrl'],
		heater.get_powers(), sim_warp_factor, config.get('Pot', {}))
	sud_task = SudTask(sud, tc, stirrer, heater, pot, DT, DT_TASK, dispatcher.msgio_get("Sud"), forecast_estimator)
	taskmgr.add(sud_task)

	# A run in progress gets force-stopped (mirrors a manual Stop) if the
	# heater or stirrer it depends on disconnects mid-brew - see SudTask.
	# check_connections()'s own comment for why.
	heater_task.set_on_connected_changed(lambda connected: sud_task.check_connections())
	stirrer_task.set_on_connected_changed(lambda connected: sud_task.check_connections())

	# How often the server/sud logs checkpoint themselves to disk (in
	# addition to always writing on their own end trigger) - so a hard
	# crash/power loss only loses up to this much data. Config-only (no CLI
	# flag): it's a plant/deployment concern like the rest of config.json,
	# not a per-invocation run-mode knob.
	log_interval = config.get('log_interval', 300)

	# Continuous server-session log - records from startup to shutdown
	# regardless of Sud state; written to logs/log_{date_time}.json on exit
	# and every log_interval seconds in between.
	server_log_task = ServerLogTask(tc, heater, heater_task, DT_TASK, args.logdir, config=config, dt=DT, log_interval=log_interval, sim_warp_factor=sim_warp_factor)
	if startup_params is not None:
		server_log_task.log_plant_params(startup_params)
	taskmgr.add(server_log_task)

	# Per-run log - only records while a Sud is actually playing; written to
	# logs/log_{date_time}_{sud_name}.log on Stop (or natural completion) and
	# every log_interval seconds in between.
	sud_log_task = SudLogTask(tc, heater, heater_task, sud, DT_TASK, args.logdir, config=config, dt=DT, log_interval=log_interval, sim_warp_factor=sim_warp_factor)
	taskmgr.add(sud_log_task)

	sud_task.set_on_plant_params(lambda params: (
		server_log_task.log_plant_params(params),
		sud_log_task.log_plant_params(params)))
	sud_task.set_on_reanchor(sud_log_task.log_reanchor)

	# 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 start: switch heater to closed-loop so the TC drives it regardless
	# of which client (browser, PyQt GUI, …) started the brew; also (re)starts
	# the per-run sud log (a Pause->resume Start is a no-op for the log).
	sud_task.set_on_start(lambda: (heater_task.start_closed_loop(), sud_log_task.start_run()))
	# Brew end (DONE/IDLE): shut heater off, broadcast open-loop/power-0 to
	# clients, and write out the per-run sud log.
	sud_task.set_on_end(lambda: (heater_task.shutdown(), sud_log_task.stop_run()))

	# Assign temp. sensor readings to tc
	sensor.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,
			'Pot': config.get('Pot', {})})
	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))

	# systemd's `stop`/`restart` send SIGTERM, whose default disposition is to
	# kill the process immediately - skipping the finally: block below
	# entirely, leaving the heater/stirrer at whatever they were last set to.
	# Routing it through the same KeyboardInterrupt path Ctrl-C (SIGINT)
	# already uses gets it the same graceful, actor-safing shutdown.
	signal.signal(signal.SIGTERM, signal.default_int_handler)

	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/
		# StirrerTask's `with device.open()`) get a chance to run before we
		# close the port - this is what actually forces the heater to 0 W
		# (AHeater.activate(False)) and the stirrer to a stopped state
		# (AStirrer.activate(False)) on any exit that reaches here.
		pending = asyncio.all_tasks(loop)
		for task in pending:
			task.cancel()
		if pending:
			loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True))
		server_log_task.write()
		# Belt-and-suspenders: explicitly deactivate both actors even if the
		# task cleanup above failed to reach their own activate(False).
		heater.close()
		stirrer.activate(False)
		http_server.shutdown()
		loop.close()
		print("Server stopped.")