import asyncio
import logging
import os
import re
import shutil
import signal
import subprocess
import sys
from pathlib import Path
from urllib.parse import urlparse

from dotenv import load_dotenv

import web
from utils.log import RawConsoleSink, color_log_line, setup_logging
from utils.vendor import check_for_updates

# run.py needs LAVALINK_* to decide whether to supervise a node. bot.py loads
# this again in its own process, which is a no-op the second time.
load_dotenv()


class _TeeStream:
    """Wraps stdout/stderr: writes go to the original stream and to the raw
    console sink. Replaces sys.stdout/sys.stderr wholesale so print(), logging,
    and traceback output all pass through it."""

    def __init__(self, original, sink):
        self._original = original
        self._sink = sink

    def write(self, s):
        n = self._original.write(s)
        self._original.flush()
        self._sink.write(s.encode("utf-8", errors="replace"))
        return n

    def flush(self):
        self._original.flush()

    def __getattr__(self, name):
        return getattr(self._original, name)


_console_sink = RawConsoleSink()  # Rolls the last run's logs/console.raw over to .1.
_real_stdout = sys.stdout  # pre-wrap, used by the bot.py output pump
sys.stdout = _TeeStream(sys.stdout, _console_sink)
sys.stderr = _TeeStream(sys.stderr, _console_sink)

# Must install the tee before setup_logging(): logging.StreamHandler() snapshots
# sys.stderr at construction, so installing the tee later would bypass it.
setup_logging(os.getenv("LOG_LEVEL", "INFO"))
logger = logging.getLogger("run")


class ProcessSupervisor:
    """Owns a child process: spawns it, watches it, and decides whether a dead
    child gets restarted (crash) or left alone (clean exit / user stop).

    Subclasses supply the command to run and how the child prefers to be stopped.
    """

    # Used in log lines, e.g. "Starting bot...".
    label = "process"

    # Windows only. True for children that install their own SIGBREAK handler and
    # shut down cleanly on it. A JVM does not: CTRL_BREAK_EVENT makes it print a
    # thread dump and keep running, so Lavalink leaves this False.
    stops_on_ctrl_break = False

    # True for children whose output is rewritten by format_line(), which needs
    # whole lines. Children that render progress with \r must leave this False,
    # since line splitting would collapse those redraws.
    reformats_output = False

    def __init__(self, sink, host_stdout):
        self.process = None
        self.status = "stopped"  # stopped | stopping | running | crashed_retrying
        self.sink = sink
        self.host_stdout = host_stdout
        self._lock = asyncio.Lock()

    def command(self) -> list[str]:
        raise NotImplementedError

    def cwd(self) -> str | None:
        return None

    def env(self) -> dict:
        return dict(os.environ)

    async def start(self):
        async with self._lock:
            if self.status in ("running", "crashed_retrying"):
                return False
            await self._spawn()
            return True

    async def stop(self):
        async with self._lock:
            proc = self.process
            if proc is None or proc.returncode is not None:
                self.status = "stopped"
                self.process = None
                return False
            # Set before _terminate, not after. Other tabs should see "stopping"
            # during the graceful shutdown instead of "stopped" too early.
            self.status = "stopping"
            await self._terminate(proc)
            self.status = "stopped"
            self.process = None
            return True

    async def restart(self):
        await self.stop()
        await self.start()

    def format_line(self, line: str) -> str | None:
        """Rewrites one line of the child's output, or returns None to drop it.
        Only called when reformats_output is set."""
        return line

    async def _spawn(self):
        creationflags = subprocess.CREATE_NEW_PROCESS_GROUP if sys.platform == "win32" else 0
        logger.info(f"Starting {self.label}...")
        self.process = await asyncio.create_subprocess_exec(
            *self.command(), creationflags=creationflags, env=self.env(), cwd=self.cwd(),
            stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.STDOUT,
            # Raised from the 64KB default so a long line (a stack trace frame,
            # a dumped payload) can still be read whole in line mode.
            limit=1_048_576,
        )
        self.status = "running"
        asyncio.create_task(self._pump_output(self.process))
        asyncio.create_task(self._watch(self.process))

    async def _terminate(self, proc):
        if sys.platform == "win32" and self.stops_on_ctrl_break:
            # The child is spawned in its own process group specifically so it can
            # be signalled independently of run.py's own console/Ctrl+C handling.
            proc.send_signal(signal.CTRL_BREAK_EVENT)
        else:
            proc.terminate()
        try:
            await asyncio.wait_for(proc.wait(), timeout=15)
        except asyncio.TimeoutError:
            proc.kill()
            await proc.wait()

    async def _pump_output(self, proc):
        """Forwards the child's merged stdout+stderr to the real host stdout and
        to the raw console sink. Reads in chunks, not lines, to preserve \\r bytes,
        unless the child's output is being reformatted."""
        try:
            if self.reformats_output:
                async for raw_line in proc.stdout:
                    line = raw_line.decode("utf-8", errors="replace").rstrip("\r\n")
                    formatted = self.format_line(line)
                    if formatted is None:
                        continue
                    self._emit(f"{formatted}\n".encode("utf-8"))
            else:
                while chunk := await proc.stdout.read(4096):
                    self._emit(chunk)
        except Exception:
            logger.exception(f"Error pumping {self.label} output")

    def _emit(self, data: bytes):
        self.host_stdout.buffer.write(data)
        self.host_stdout.buffer.flush()
        self.sink.write(data)

    async def _watch(self, proc):
        returncode = await proc.wait()
        if self.process is not proc:
            # Superseded by a newer spawn, e.g. stop+start raced the watcher. Ignore.
            return

        if self.status in ("stopped", "stopping"):
            # stop() is already handling this same proc via the same proc.wait().
            # Let it own the status transition instead of racing it here.
            self.process = None
            return

        if returncode == 0:
            logger.info(f"{self.label.capitalize()} stopped cleanly (exit code 0). Not restarting.")
            self.status = "stopped"
            self.process = None
            return

        logger.warning(f"{self.label.capitalize()} stopped (exit code {returncode}). Restarting in 5 seconds...")
        self.status = "crashed_retrying"
        self.process = None
        await asyncio.sleep(5)
        async with self._lock:
            if self.status == "crashed_retrying":
                await self._spawn()


class BotSupervisor(ProcessSupervisor):
    label = "bot"
    stops_on_ctrl_break = True

    def command(self):
        return [sys.executable, "bot.py"]

    def env(self):
        # Piped (not a real console) means PEP 528's UTF-8-for-console default no
        # longer applies; force UTF-8 explicitly to match what the sink expects.
        return {**os.environ, "PYTHONIOENCODING": "utf-8"}


# Spring Boot's default console line, e.g.
# 2026-09-06T15:19:56.714+01:00  INFO 6156 --- [Lavalink] [   main] o.s.b.w.e.u.UndertowWebServer : ...
# The PID, the "---" separator and the app/thread brackets are all dropped in the
# rewrite; only the timestamp, level, logger and message survive.
_SPRING_LINE_RE = re.compile(
    r"^(?P<year>\d{4})-(?P<month>\d{2})-(?P<day>\d{2})T(?P<time>\d{2}:\d{2}:\d{2})\.\d+\S*\s+"
    r"(?P<level>[A-Z]+)\s+\d+\s+---\s+(?:\[[^\]]*\]\s*)+(?P<logger>\S+)\s*:\s(?P<msg>.*)$"
)

# Node loggers whose every line restates something the bot already reports, or
# that carry no signal at all. Dropped from the console along with the trace body
# that follows them; the node's own logs/ files are untouched either way.
#
# RequestLoggingFilter logs one line per HTTP call with the full payload inlined,
# which for a play is an encoded track blob plus all fifteen equalizer bands.
# LocalAudioTrackExecutor logs playback failures ("Suspicious exception for
# playback of X", "Error in playback of X") as a Java cause chain, one stack
# trace per YouTube client tried, which MusicManager reports as a single line
# naming the track and the stand-in it queued.
_QUIET_NODE_LOGGERS = frozenset({
    "lavalink.RequestLoggingFilter",
    "lavalink.LocalAudioTrackExecutor",
})

# A stack frame in a Java trace: "\tat pkg.Class.method(File.java:1)" or
# "\t... 65 common frames omitted". Dropped under a WARNING or ERROR from any
# other node logger, which keeps the exception and "Caused by:" lines. A failed
# load prints about 80 of these, most from the HTTP server libraries the node
# is built on (Undertow, Spring) rather than its audio code.
_TRACE_FRAME_RE = re.compile(r"^\s+(?:at |\.\.\. \d+ )")

# Rewrites applied to a node line's message before it reaches the console. The
# line is kept, only the matched value is replaced; the node's own logs/ files
# are untouched.
#
# YoutubeAccessTokenTracker prints the visitor id it obtained in full, which
# identifies the node's YouTube session.
_REDACTIONS = (
    (re.compile(r"(visitor id succeeded, new one is )[^\s,]+"), r"\1<redacted>"),
)


def _redact(msg: str) -> str:
    for pattern, replacement in _REDACTIONS:
        msg = pattern.sub(replacement, msg)
    return msg


class LavalinkSupervisor(ProcessSupervisor):
    """Runs the Lavalink node the music cogs connect to.

    Optional: with no LAVALINK_DIR configured, or the jar not installed there,
    the node is simply never started and the music commands report themselves as
    unavailable. The rest of the bot is unaffected.
    """

    label = "lavalink"
    reformats_output = True

    def __init__(self, sink, host_stdout):
        super().__init__(sink, host_stdout)
        directory = os.getenv("LAVALINK_DIR")
        self.directory = Path(directory).expanduser() if directory else None

        # Whether the last prefixed line was dropped, so its unprefixed
        # continuation lines can be dropped with it. See format_line().
        self._dropping = False
        # Whether the last prefixed line was a WARNING or ERROR, so the stack
        # frames under it can be dropped.
        self._trimming = False

        # The readiness probe targets whatever LAVALINK_URI points at, so changing
        # the port in application.yml and .env doesn't leave it polling the old one.
        parsed = urlparse(os.getenv("LAVALINK_URI") or "http://127.0.0.1:2333")
        self.host = parsed.hostname or "127.0.0.1"
        self.port = parsed.port or 2333

    @property
    def jar(self) -> Path | None:
        return self.directory / "Lavalink.jar" if self.directory else None

    def unavailable_reason(self) -> str | None:
        """Why the node can't be started, or None if it can."""
        if not self.directory:
            return "LAVALINK_DIR is not set"
        if not self.jar.is_file():
            return f"no Lavalink.jar in {self.directory}"
        if not (self.directory / "application.yml").is_file():
            return f"no application.yml in {self.directory}"
        if not shutil.which("java"):
            return "java is not on PATH"
        return None

    def command(self):
        # -Xmx caps the heap so a long uptime can't let the JVM balloon on a box
        # that is also running the bot and the dashboard.
        return [shutil.which("java"), "-Xmx512M", "-jar", str(self.jar)]

    def cwd(self):
        # Lavalink resolves application.yml, plugins/ and logs/ relative to cwd.
        return str(self.directory)

    def format_line(self, line):
        """Restates a Spring Boot log line in the bot's own console format, so the
        node's output reads as part of the same log instead of a second one pasted
        in. Anything that doesn't match (the startup banner, a trace's exception
        lines) is left exactly as printed, and _REDACTIONS is applied to whatever
        does match. Returns None for a line dropped as noise."""
        match = _SPRING_LINE_RE.match(line)
        if not match:
            # A stack trace line or a blank line inside one. These carry no
            # prefix of their own, so they belong to whichever prefixed line came
            # last and are kept or dropped with it.
            if self._dropping:
                return None
            # Under a kept WARNING or ERROR, only the frames and blanks go.
            if self._trimming and (not line.strip() or _TRACE_FRAME_RE.match(line)):
                return None
            return line

        # WARN -> WARNING keeps the level names, and so the colors, in step with
        # logging's own. The logger is cut to its last segment and namespaced, so
        # o.s.b.w.e.u.UndertowWebServer reads as lavalink.UndertowWebServer.
        level = "WARNING" if match["level"] == "WARN" else match["level"]
        timestamp = f"{match['day']}/{match['month']}/{match['year']} {match['time']}"
        name = f"lavalink.{match['logger'].rsplit('.', 1)[-1]}"

        # At DEBUG the whole cause chain is the point, so nothing is dropped.
        debug = logger.isEnabledFor(logging.DEBUG)
        self._dropping = name in _QUIET_NODE_LOGGERS and not debug
        self._trimming = level in ("WARNING", "ERROR") and not debug
        if self._dropping:
            return None
        return color_log_line(timestamp, level, name, _redact(match["msg"]))

    async def wait_until_ready(self, timeout: float = 90):
        """Blocks until the node accepts connections, so the bot doesn't come up
        swinging at a port that isn't listening yet. Returns whether it did."""
        deadline = asyncio.get_running_loop().time() + timeout
        while asyncio.get_running_loop().time() < deadline:
            if self.status not in ("running", "crashed_retrying"):
                return False
            try:
                _, writer = await asyncio.open_connection(self.host, self.port)
                writer.close()
                await writer.wait_closed()
                logger.info("Lavalink is accepting connections")
                return True
            except OSError:
                await asyncio.sleep(1)
        logger.warning(f"Lavalink did not become ready within {timeout:.0f}s, starting the bot anyway")
        return False


class WebState:
    """Hot-reload bookkeeping for web.py's uvicorn server, kept outside the
    web module itself since /web/reload replaces the module via importlib.reload."""

    def __init__(self):
        self.web_server = None
        self.web_epoch = 0
        # Serializes /web/reload attempts so a double-click can't race two
        # concurrent reloads against each other and each other's server refs.
        self.reload_lock = asyncio.Lock()
        # Bumped on every cog load/unload/reload so other tabs notice and refresh.
        self.cogs_epoch = 0
        # Transient badges are stored as deadlines, not rendered flags, so their
        # lifetime survives the surrounding region being re-rendered.
        self.cog_badges = {}  # extension -> (deadline, error or None)
        self.last_sync = None  # (deadline, count, error)
        self.started_deadline = None
        self.bot_was_ready = None  # None until first observed
        # Set by main(), not read yet. Groundwork for a Lavalink status row on the
        # dashboard alongside the bot's. Carried on the state rather than added to
        # web.start()'s signature, which /web/reload also calls.
        self.lavalink = None


async def main():
    supervisor = BotSupervisor(_console_sink, _real_stdout)
    lavalink = LavalinkSupervisor(_console_sink, _real_stdout)
    web_state = WebState()
    web_state.lavalink = lavalink

    # Started before the bot so the node is listening by the time the music cog
    # tries to connect. A missing or unconfigured install is not fatal.
    reason = lavalink.unavailable_reason()
    if reason:
        logger.warning(f"Lavalink not started ({reason}). Music commands will be unavailable.")
    else:
        await lavalink.start()
        await lavalink.wait_until_ready()

    await supervisor.start()

    # Backgrounded, not awaited directly: /web/reload tears this task down and starts
    # a fresh one in its place, and that must not end run.py's own lifetime.
    web_task = asyncio.create_task(web.start(supervisor, web_state))
    web_task.add_done_callback(web._log_task_exception)

    # Reports newer releases of the dashboard's vendored libraries. Backgrounded
    # and failure-tolerant: it writes nothing, and nothing here waits on it. The
    # notice lands in the console page because this process's stdout is teed there.
    # Bound to a local because asyncio only holds a weak reference to a running
    # task, and this one has no other owner to keep it from being collected.
    update_check = asyncio.create_task(check_for_updates())

    try:
        await asyncio.Future()  # run until interrupted (Ctrl+C / CTRL_BREAK_EVENT)
    finally:
        # Bot first: it should leave its voice channels before the node it is
        # streaming through goes away.
        await supervisor.stop()
        await lavalink.stop()
        _console_sink.close()


if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        logger.info("Stopped.")