from .Challenge import Challenge

import subprocess
import time
import re
import os

class Carbeat(Challenge):
    flag_location    = 'flags/carbeat.txt'
    history_location = 'history/carbeat.txt'

    _CONTAINER   = "carbeat_container"
    _SERVICE_CMD = ["docker", "exec", "-i", _CONTAINER, "/home/ctfuser/chall/mybini"]
    _HEX_RE      = re.compile(r'^[0-9a-fA-F]+$')

    def _read_container_flag(self) -> str:
        out = subprocess.run(["docker", "exec", self._CONTAINER, "cat", "/flag.txt"],
                             capture_output=True, text=True)
        if out.returncode != 0 or not out.stdout.strip():
            raise FileNotFoundError("Flag not found in container (/flag.txt)")
        return out.stdout.strip()

    def _spawn(self):
        return subprocess.Popen(
            self._SERVICE_CMD,
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
            bufsize=0,
        )

    def _read_until(self, proc, token, timeout=5.0, max_bytes=1_000_000):
        start = time.time()
        buf = []
        r = proc.stdout.read
        while True:
            if time.time() - start > timeout:
                tail = ''.join(buf)[-500:]
                raise TimeoutError(f"Timeout waiting for '{token}'. Got so far:\n{tail}")
            ch = r(1)
            if ch == "" and proc.poll() is not None:
                raise RuntimeError(f"Process ended while waiting for '{token}'. Output:\n{''.join(buf)}")
            buf.append(ch)
            if len(buf) > max_bytes:
                raise RuntimeError("Exceeded max read size")
            if token in "".join(buf):
                return "".join(buf)

    def _send_line(self, proc, s: str):
        proc.stdin.write(s + "\n")
        proc.stdin.flush()

    def _expect_hex_field(self, text: str, label: str) -> str:
        m = re.search(rf"{re.escape(label)}\s*:\s*([0-9a-fA-F]+)", text)
        assert m, f"Missing '{label}' in output. Tail:\n{text[-400:]}"
        hx = m.group(1)
        assert self._HEX_RE.match(hx), f"{label} is not hex"
        return hx

    def distribute(self, flag):
        try:
            os.makedirs(os.path.dirname(self.flag_location), exist_ok=True)
            with open(self.flag_location, 'w') as f:
                f.write(flag)

            os.makedirs(os.path.dirname(self.history_location), exist_ok=True)
            with open(self.history_location, 'a') as f:
                f.write(flag + '\n')

            self.logger.info(f'Flag {flag} written to {self.flag_location}')
            return True
        except Exception as e:
            self.logger.error(f'Could not write flag to {self.flag_location}: {e}')
            return False

                
    def register(self, proc, name, level, timeout=1.):
        self._read_until(proc, "> ", timeout=timeout)
        self._send_line(proc, "1")
        self._read_until(proc, ": ", timeout=timeout)
        self._send_line(proc, name)
        self._read_until(proc, ": ", timeout=timeout)
        self._send_line(proc, str(level))

    def login(self, proc, name, timeout=1.):
        self._read_until(proc, "> ", timeout=timeout)
        self._send_line(proc, "2")
        self._read_until(proc, ": ", timeout=timeout)
        self._send_line(proc, name)

    def list_user(self, proc, timeout=1.):
        self._read_until(proc, "> ", timeout=timeout)
        self._send_line(proc, "3")
        return self._read_until(proc, "\n1. register")
    
    def bini(self, proc, id, data, timeout=1.0):
        self._read_until(proc, "> ", timeout=timeout)
        self._send_line(proc, str(id))
        self._read_until(proc, ": ", timeout=timeout)
        self._send_line(proc, data)

    def logout(self, proc, timeout=1.):
        self._read_until(proc, "> ", timeout=timeout)
        self._send_line(proc, "4")

    def exit(self, proc, timeout=1.):
        self._read_until(proc, "> ", timeout=timeout)
        self._send_line(proc, "3")

    def check(self):
        try:
            with open(self.flag_location, 'r') as f:
                host_flag = f.read().strip()
            container_flag = self._read_container_flag()
            assert host_flag == container_flag, 'Flag mismatch between host and container'
            self.logger.info('[ok] flag parity (carbeat)')

            proc = self._spawn()

            self.register(proc, "karbit", 1337)
            
            name = self.list_user(proc)
            assert "karbit" in name
            
            self.login(proc, "karbit")
            self.bini(proc, 1, "waguri")
            self.bini(proc, 2, "0")
            waguri = self._read_until(proc, "0. edit", timeout=1.)
            assert "waguri" in waguri
            
            self.bini(proc, 3, "0")
            self.logout(proc)
            self.exit(proc)
            
            try:
                proc.wait(timeout=2.0)
            except subprocess.TimeoutExpired:
                proc.kill()
                raise AssertionError("Program did not exit after option 3")
            self.logger.info("[ok] service exit on 3")

            proc_alarm = self._spawn()
            self._read_until(proc_alarm, "zzz", timeout=190.0)  # 180s + slack
            try:
                proc_alarm.wait(timeout=5.0)
            except subprocess.TimeoutExpired:
                proc_alarm.kill()
                raise AssertionError("Alarm fired but process did not exit")
            self.logger.info("[ok] alarm fired ('zzz') and process self-terminated")

            self.logger.info('Check passed for carbeat')
            return True

        except Exception as e:
            self.logger.error(f'Could not check carbeat: {e}')
            return False