"""Прокси с heartbeat для штатного интерфейса MuScriptor.

Зачем: /transcribe отдаёт SSE-поток, который на бесплатном туннеле рвётся
по простою (провайдер закрывает соединение, если данные не идут ~60 с).
Прокси читает поток с сервера отдельным потоком и, пока он молчит,
отправляет браузеру комментарий-пульс ": ping", поэтому туннель видит трафик
и не рвёт соединение. Интерфейс остаётся штатным, в пакет не лезем.

Запуск: python plus/sse_proxy.py --listen 8223 --target 8222
"""

import argparse
import http.client
import json
import os
import re
import queue
import sys
import threading
import time
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

MIDI_TYPES = ("audio/midi", "audio/x-midi", "application/x-midi",
              "application/octet-stream", "application/midi")


def guess_audio(body: bytes, content_type: str) -> tuple[str, bytes] | None:
    """Вытащить аудио из multipart-запроса, чтобы потом мерить по нему громкость."""
    marker = "boundary="
    if "multipart/form-data" not in content_type or marker not in content_type:
        return None
    boundary = content_type.split(marker, 1)[1].split(";")[0].strip().strip('"')
    parts = body.split(b"--" + boundary.encode())
    for part in parts:
        if b"filename=" not in part:
            continue
        head, _, tail = part.partition(b"\r\n\r\n")
        if not tail:
            continue
        name = ""
        for line in head.decode("latin-1").split("\r\n"):
            if "filename=" in line:
                found = re.search(r'filename="([^"]+)"', line)
                if found:
                    name = found.group(1)
                else:
                    name = line.split("filename=")[-1].strip().strip('"').split(";")[0]
        suffix = ("." + name.rsplit(".", 1)[-1].lower()) if "." in name else ""
        if suffix in (".wav", ".mp3", ".flac", ".ogg", ".m4a", ".aiff", ".aif", ".opus"):
            data = tail.rstrip(b"\r\n")
            data = data[:-2] if data.endswith(b"\r\n") else data
            return name or ("audio" + suffix), data
    return None


def apply_velocity(audio: tuple[str, bytes] | None, midi: bytes) -> tuple[bytes, dict]:
    """Проставить измеренную громкость вместо константы 100."""
    if not audio:
        return midi, {"skipped": "нет аудио для замера"}
    import tempfile
    import velocity as vel
    name, blob = audio
    suffix = ("." + name.rsplit(".", 1)[-1].lower()) if "." in name else ".wav"
    try:
        with tempfile.NamedTemporaryFile(suffix=suffix, delete=True) as handle:
            handle.write(blob)
            handle.flush()
            return vel.apply(handle.name, midi)
    except Exception as problem:
        return midi, {"skipped": "замер не вышел: %s" % problem}


sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import fusion  # noqa: E402
import verify_notes  # noqa: E402
import articulation  # noqa: E402
WIDE_INSTRUMENTS = ("acoustic_piano,electric_piano,chromatic_percussion,organ,acoustic_guitar,clean_electric_guitar,acoustic_bass,electric_bass,string_ensemble,voice,synth_pad,drums")

ALL_INSTRUMENTS = ("acoustic_piano,electric_piano,chromatic_percussion,organ,acoustic_guitar,clean_electric_guitar,distorted_electric_guitar,acoustic_bass,electric_bass,violin,viola,cello,contrabass,orchestral_harp,timpani,string_ensemble,synth_strings,voice,orchestra_hit,trumpet,trombone,tuba,french_horn,brass_section,soprano_and_alto_sax,tenor_sax,baritone_sax,oboe,english_horn,bassoon,clarinet,flutes,synth_lead,synth_pad,drums")


def expand_instruments(text: str) -> list[str]:
    """Разбор набора: all → все группы, auto → пусто, иначе список через запятую."""
    cleaned = (text or "").strip()
    if not cleaned or cleaned.lower() == "auto":
        return []
    if cleaned.lower() in ("all", "все"):
        return ALL_INSTRUMENTS.split(",")
    if cleaned.lower() in ("wide", "широкий"):
        return WIDE_INSTRUMENTS.split(",")
    return [name.strip() for name in cleaned.split(",") if name.strip()]

MAX_GROUPS = 10                      # шире модель размазывается: замерено
BAND_LIST = "drums,electric_bass,acoustic_bass,chromatic_percussion,acoustic_piano,electric_piano,clean_electric_guitar,synth_pad".split(",")


def guard_instruments(names: list[str], label: str = "") -> tuple[list[str], str]:
    """Свести опасный набор к безопасному и объяснить почему."""
    if not names:
        return names, ""
    band_alias = [n for n in names if n.lower() in ("band", "набор")]
    if band_alias:
        return list(BAND_LIST), "%s: пресет band, беру 8 проверенных групп" % label
    known = set(ALL_INSTRUMENTS.split(","))
    unknown = [n for n in names if n not in known]
    if unknown:
        names = [n for n in names if n in known]
        if not names:
            return list(BAND_LIST), ("%s: неизвестные группы %s, беру band"
                                     % (label, ", ".join(unknown[:4])))
        return names, "%s: выкинул неизвестные группы %s" % (label, ", ".join(unknown[:4]))
    if len(names) > MAX_GROUPS:
        return list(BAND_LIST), (
            "%s: %d групп — модель размазывается (замер: 35 групп дали 0 нот,"
            " 12 групп хуже). Беру проверенный band из %d групп."
            % (label or "набор", len(names), len(BAND_LIST)))
    return names, ""

def replace_instruments(body: bytes, content_type: str, names: list[str]) -> bytes:
    """Заменить поля instruments на нужный набор (для второго прохода)."""
    marker = "boundary="
    if "multipart/form-data" not in content_type or marker not in content_type:
        return body
    boundary = content_type.split(marker, 1)[1].split(";")[0].strip().strip('"')
    chunks = body.split(b"--" + boundary.encode())
    kept = []
    for chunk in chunks:
        if b'name="instruments"' in chunk:
            continue
        kept.append(chunk)
    rebuilt = (b"--" + boundary.encode()).join(kept)
    if names:
        rebuilt = add_instruments(rebuilt, content_type, names)
    return rebuilt


def run_pass(host: str, port: int, path: str, body: bytes, content_type: str,
             instruments: list[str]) -> str:
    """Один проход транскрипции к серверу, ответ как текст (SSE)."""
    payload = replace_instruments(body, content_type, instruments)
    connection = http.client.HTTPConnection(host, port, timeout=3600)
    headers = {"Content-Type": content_type, "Content-Length": str(len(payload)),
               "Accept-Encoding": "identity"}
    connection.request("POST", path, body=payload, headers=headers)
    answer = connection.getresponse()
    text = answer.read().decode("utf-8", "ignore")
    connection.close()
    return text


def build_midi(host: str, port: int, notes: list[dict]) -> bytes:
    """Собрать MIDI из слитых нот силами самого сервера (/export)."""
    payload = json.dumps(fusion.notes_to_export(notes)).encode("utf-8")
    connection = http.client.HTTPConnection(host, port, timeout=600)
    connection.request("POST", "/export", body=payload,
                       headers={"Content-Type": "application/json"})
    answer = connection.getresponse()
    data = answer.read()
    connection.close()
    return data


WARNING_BYPASS = """<script>/*serveo-warning-bypass*/
(function(){var H="serveo-skip-browser-warning",V="true";
try{var f=window.fetch;
if(f){window.fetch=function(i,o){try{
if(typeof i==="string"||i instanceof URL){o=o||{};var h=new Headers(o.headers||{});h.set(H,V);o.headers=h;}
else if(i instanceof Request){var g=new Headers(i.headers||{});g.set(H,V);i=new Request(i,{headers:g});}
}catch(e){}return f.call(this,i,o);};}
var sd=XMLHttpRequest.prototype.send;
XMLHttpRequest.prototype.send=function(){try{this.setRequestHeader(H,V);}catch(e){}
return sd.apply(this,arguments);};}catch(e){}})();
</script>"""


def inject_warning_bypass(data: bytes) -> tuple[bytes, bool]:
    """Вставить обход предупреждения serveo в HTML интерфейса."""
    if os.environ.get("PX_NO_INJECT") == "1" or b"serveo-warning-bypass" in data:
        return data, False
    low = data.lower()
    for needle in (b"</head>", b"<body"):
        at = low.find(needle)
        if at >= 0:
            return data[:at] + WARNING_BYPASS.encode("utf-8") + data[at:], True
    return data, False


HOP_HEADERS = {"connection", "keep-alive", "proxy-authenticate", "proxy-authorization",
               "te", "trailers", "transfer-encoding", "upgrade", "host",
               "content-length", "accept-encoding"}
PING_SECONDS = 10.0
# Звуковой банк (38 МБ) не гоняем через туннель: браузер берёт его с CDN.
FONT_CDN = ("https://huggingface.co/Nokia123543/muscriptor-large-fork"
            "/resolve/main/soundfonts/MuseScore_General.sf3")


def read_body(handler) -> bytes:
    length = handler.headers.get("Content-Length")
    if length:
        return handler.rfile.read(int(length))
    if (handler.headers.get("Transfer-Encoding") or "").lower() == "chunked":
        parts = []
        while True:
            line = handler.rfile.readline().strip()
            size = int(line.split(b";")[0] or b"0", 16)
            if size == 0:
                handler.rfile.readline()
                break
            parts.append(handler.rfile.read(size))
            handler.rfile.readline()
        return b"".join(parts)
    return b""


BAND = ("drums,electric_bass,acoustic_bass,chromatic_percussion,acoustic_piano,"
        "electric_piano,clean_electric_guitar,distorted_electric_guitar,"
        "string_ensemble,synth_pad,synth_lead,voice")


def add_instruments(body: bytes, content_type: str, names: list[str]) -> bytes:
    """Дописать поля instruments в multipart, если их там нет."""
    if not names or b'name="instruments"' in body:
        return body
    marker = "boundary="
    if marker not in content_type:
        return body
    boundary = content_type.split(marker, 1)[1].split(";")[0].strip().strip('"')
    closing = b"--" + boundary.encode() + b"--"
    if closing not in body:
        return body
    parts = b""
    for name in names:
        parts += (b"--" + boundary.encode() + b"\r\n"
                  b'Content-Disposition: form-data; name="instruments"\r\n\r\n'
                  + name.encode() + b"\r\n")
    return body.replace(closing, parts + closing, 1)


POST_ONLY = ("/transcribe", "/transcribe/midi", "/export", "/auralize", "/sheets")


def serve_alias(handler, target: str, method: str = "GET") -> bool:
    """Отдать другой путь с того же сервера (например favicon)."""
    try:
        upstream = http.client.HTTPConnection(handler.target_host, handler.target_port,
                                              timeout=30)
        upstream.request(method, target, headers={"Accept-Encoding": "identity"})
        answer = upstream.getresponse()
        data = answer.read()
    except Exception as problem:
        print("алиас %s не вышел: %s" % (target, problem), flush=True)
        return False
    handler.send_response(answer.status)
    for key, value in answer.getheaders():
        if key.lower() in HOP_HEADERS or key.lower() == "content-length":
            continue
        handler.send_header(key, value)
    handler.send_header("Content-Length", str(len(data)))
    handler.end_headers()
    if method != "HEAD":
        handler.wfile.write(data)
    upstream.close()
    return True


class Proxy(BaseHTTPRequestHandler):
    protocol_version = "HTTP/1.1"
    target_host = "127.0.0.1"
    target_port = 8222
    instruments: list[str] = []
    measure_velocity = True
    ping_seconds = PING_SECONDS
    multi_pass: list[str] = []
    energy_threshold = 6.0
    refine_articulation = True
    last_audio: tuple[str, bytes] | None = None
    font_cdn = os.environ.get("MUSCRIPTOR_FONT_CDN", FONT_CDN)

    def log_message(self, fmt, *args):
        pass

    def do_GET(self):
        self.pass_through("GET")

    def do_HEAD(self):
        self.pass_through("HEAD")

    def do_POST(self):
        self.pass_through("POST")

    def do_PUT(self):
        self.pass_through("PUT")

    def do_DELETE(self):
        self.pass_through("DELETE")

    def do_OPTIONS(self):
        self.pass_through("OPTIONS")

    def pass_through(self, method: str):
        path_only = self.path.split("?")[0]
        self._started = time.time()
        declared = self.headers.get("Content-Length") or "-"
        print("ЗАПРОС %s %s | %s байт | клиент %s"
              % (method, path_only, declared, self.client_address[0]), flush=True)
        if (method in ("GET", "HEAD") and self.font_cdn
                and path_only.startswith("/soundfonts/")
                and "local=1" not in self.path):
            # тяжёлый банк отдаёт CDN, туннель отдыхает
            print("банк: редирект на CDN (%s)" % path_only, flush=True)
            self.send_response(302)
            self.send_header("Location", self.font_cdn)
            self.send_header("Content-Length", "0")
            self.send_header("Cache-Control", "public, max-age=86400")
            self.end_headers()
            return
        # 1. браузер всегда просит иконку: отдаём png вместо JSON-404
        if path_only == "/favicon.ico" and method in ("GET", "HEAD"):
            if serve_alias(self, "/favicon.png", method):
                return
        # 2. GET на POST-маршрут: понятный текст вместо {"detail":"Not Found"}
        if method in ("GET", "HEAD") and path_only in POST_ONLY:
            note = ("Этот адрес принимает только POST.\n"
                    "Открой интерфейс: %s/\n" % self.headers.get("Host", "ссылка"))
            payload = note.encode("utf-8")
            self.send_response(405)
            self.send_header("Content-Type", "text/plain; charset=utf-8")
            self.send_header("Allow", "POST")
            self.send_header("Content-Length", str(len(payload)))
            self.end_headers()
            if method != "HEAD":
                self.wfile.write(payload)
            print("ПОДСКАЗКА %s %s | 405 текстом" % (method, path_only), flush=True)
            return
        # 3. лишний слэш в конце: /transcribe/ → /transcribe
        if method == "POST" and path_only != "/" and path_only.endswith("/"):
            trimmed = path_only.rstrip("/")
            if trimmed in POST_ONLY:
                query = self.path[len(path_only):]
                self.path = trimmed + query
                path_only = trimmed

        body = read_body(self)
        content_type = self.headers.get("Content-Type") or ""
        audio = guess_audio(body, content_type)
        if os.environ.get("PXDEBUG"):
            print("PXDEBUG path=%s body=%d ctype=%s filename=%s audio=%s"
                  % (self.path, len(body), content_type[:60],
                     b"filename=" in body, bool(audio)), flush=True)
        if audio:
            Proxy.last_audio = audio
        if self.quality_applies(method):
            self.serve_quality(body)
            return
        if (self.instruments and method == "POST"
                and self.path.split("?")[0] in ("/transcribe", "/transcribe/midi")):
            body = add_instruments(body, self.headers.get("Content-Type") or "",
                                   self.instruments)
        upstream = http.client.HTTPConnection(self.target_host, self.target_port, timeout=900)
        headers = {k: v for k, v in self.headers.items()
                   if k.lower() not in HOP_HEADERS}
        headers["Accept-Encoding"] = "identity"
        try:
            upstream.request(method, self.path, body=body, headers=headers)
            answer = upstream.getresponse()
        except Exception as problem:
            self.send_error(502, "upstream: %s" % problem)
            return

        content_type = (answer.getheader("Content-Type") or "")
        self.send_response(answer.status)
        stream = content_type.startswith("text/event-stream")
        for key, value in answer.getheaders():
            low = key.lower()
            if low in HOP_HEADERS or low in ("server", "date"):
                continue
            if stream and low in ("cache-control", "x-accel-buffering"):
                continue
            self.send_header(key, value)
        if stream:
            self.begin_stream()
            try:
                self.stream_events(answer)
            finally:
                self.end_stream()
        else:
            data = answer.read()
            content_type = (answer.getheader("Content-Type") or "").lower()
            if "text/html" in content_type:
                data, injected = inject_warning_bypass(data)
                if injected:
                    print("интерфейс: добавлен обход предупреждения провайдера", flush=True)
            looks_midi = content_type.startswith(MIDI_TYPES) or data[:4] == b"MThd"
            if (self.measure_velocity and method != "HEAD" and looks_midi
                    and data[:4] == b"MThd"):
                data, stats = apply_velocity(self.last_audio, data)
                print("velocity: %s" % stats, flush=True)
            if method == "HEAD":
                length = answer.getheader("Content-Length")
                self.send_header("Content-Length", length if length is not None else "0")
            elif answer.status not in (204, 304) and answer.status >= 200:
                self.send_header("Content-Length", str(len(data)))
            self.end_headers()
            if method != "HEAD" and answer.status not in (204, 304):
                self.wfile.write(data)
            if path_only != "/health":
                print("ОТВЕТ %s | статус %d | %d байт | %.2f с"
                      % (path_only, answer.status, len(data),
                         time.time() - getattr(self, "_started", time.time())), flush=True)
        upstream.close()

    def begin_stream(self) -> None:
        """Заголовки SSE: chunked + Connection: close, чтобы туннель (HTTP/2)
        видел конец ответа и не рвал поток с ERR_HTTP2_PROTOCOL_ERROR."""
        self.send_header("Cache-Control", "no-cache, no-transform")
        self.send_header("X-Accel-Buffering", "no")
        self.send_header("Transfer-Encoding", "chunked")
        self.send_header("Connection", "close")
        self.end_headers()
        self.close_connection = True
        self.emit(b": open\n\n")

    def emit(self, data: bytes) -> None:
        if not data:
            return
        self.wfile.write(b"%x\r\n" % len(data) + data + b"\r\n")
        self.wfile.flush()

    def end_stream(self) -> None:
        self.close_connection = True
        try:
            self.wfile.write(b"0\r\n\r\n")
            self.wfile.flush()
        except Exception:
            pass

    def refine(self, notes: list[dict], skip_filter: bool = False) -> list[dict]:
        """Пересчитать концы нот и артикуляцию по записи."""
        if not notes or not self.last_audio:
            return notes
        import tempfile
        name, blob = self.last_audio
        suffix = ("." + name.rsplit(".", 1)[-1].lower()) if "." in name else ".wav"
        try:
            with tempfile.NamedTemporaryFile(suffix=suffix, delete=True) as handle:
                handle.write(blob)
                handle.flush()
                refined, stats = articulation.refine([dict(n) for n in notes], handle.name)
        except Exception as problem:
            print("артикуляция не вышла: %s" % problem, flush=True)
            return notes
        print("артикуляция: %s" % stats, flush=True)
        for note in refined:
            note.pop("measured_end", None)
        return self.filter_by_audio(refined) if skip_filter else refined

    def filter_by_audio(self, notes: list[dict]) -> list[dict]:
        """Оставить только то, что подтверждается энергией в записи."""
        if not notes or not self.last_audio or self.energy_threshold <= 0:
            return notes
        import tempfile
        name, blob = self.last_audio
        suffix = ("." + name.rsplit(".", 1)[-1].lower()) if "." in name else ".wav"
        try:
            with tempfile.NamedTemporaryFile(suffix=suffix, delete=True) as handle:
                handle.write(blob)
                handle.flush()
                kept, dropped, stats = verify_notes.verify(
                    notes, handle.name, threshold=self.energy_threshold)
        except Exception as problem:
            print("проверка по аудио не вышла: %s" % problem, flush=True)
            return notes
        print("проверка по аудио: %s" % stats, flush=True)
        for note in kept:
            note.pop("energy_ratio", None)
        return kept

    def quality_applies(self, method: str) -> bool:
        """Нужен ли конвейер качества для этого запроса."""
        if method != "POST" or not self.last_audio:
            return False
        if self.path.split("?")[0] != "/transcribe":
            return False
        return bool(self.multi_pass) or self.refine_articulation

    def serve_quality(self, body: bytes) -> None:
        """Ответ потоком с пульсом сразу, конвейер считается в фоне."""
        outcome: queue.Queue = queue.Queue()

        def worker():
            try:
                text = run_pass(self.target_host, self.target_port, self.path, body,
                                self.headers.get("Content-Type") or "", self.instruments)
                notes = fusion.parse_sse(text)
                if self.multi_pass:
                    second = run_pass(self.target_host, self.target_port, self.path, body,
                                      self.headers.get("Content-Type") or "", self.multi_pass)
                    notes, stats = fusion.merge([notes, fusion.parse_sse(second)])
                    print("слияние: %s" % stats, flush=True)
                    notes = self.filter_by_audio(notes)
                    if self.refine_articulation:
                        notes = self.refine(notes, skip_filter=True)
                elif self.refine_articulation:
                    notes = self.refine(notes)
                outcome.put(fusion.to_sse(notes))
            except Exception as problem:
                print("конвейер упал: %s" % problem, flush=True)
                outcome.put(None)

        threading.Thread(target=worker, daemon=True).start()
        self.send_response(200)
        self.send_header("Content-Type", "text/event-stream")
        try:
            self.begin_stream()
        except Exception:
            return
        payload = None
        try:
            while True:
                try:
                    payload = outcome.get(timeout=self.ping_seconds)
                    break
                except queue.Empty:
                    self.emit(b": ping\n\n")
            if payload is None:
                payload = fusion.to_sse([])
            print("ОТВЕТ %s | конвейер %.1f с | %d байт"
                  % (self.path.split("?")[0],
                     time.time() - getattr(self, "_started", time.time()),
                     len(payload)), flush=True)
            self.emit(payload)
            self.emit(b": done\n\n")
        except Exception:
            pass
        finally:
            self.end_stream()

    def stream_events(self, answer):
        """Переливаем SSE, вставляя ": ping" во время молчания модели."""
        lines: queue.Queue = queue.Queue()

        def pump():
            try:
                while True:
                    line = answer.readline()
                    if not line:
                        break
                    lines.put(line)
            except Exception:
                pass
            finally:
                lines.put(None)

        threading.Thread(target=pump, daemon=True).start()
        while True:
            try:
                line = lines.get(timeout=PING_SECONDS)
            except queue.Empty:
                try:
                    self.emit(b": ping\n\n")
                except Exception:
                    return
                continue
            if line is None:
                return
            try:
                self.emit(line)
            except Exception:
                return


def main():
    parser = argparse.ArgumentParser(description="SSE-прокси MuScriptor с heartbeat")
    parser.add_argument("--listen", type=int, default=8223)
    parser.add_argument("--target", type=int, default=8222)
    parser.add_argument("--ping", type=float, default=PING_SECONDS)
    parser.add_argument("--font-cdn", default=os.environ.get("MUSCRIPTOR_FONT_CDN", FONT_CDN),
                        help="адрес звукового банка на CDN; none — отдавать с сервера")
    parser.add_argument("--multi-pass", default="",
                        help="второй проход с этим набором инструментов и слияние,"
                             " например: chromatic_percussion,electric_bass,acoustic_bass")
    parser.add_argument("--no-articulation", action="store_true",
                        help="не пересчитывать концы нот и не ставить легато/стакато")
    parser.add_argument("--energy-threshold", type=float, default=6.0,
                        help="во сколько раз тон должен превышать соседние полутона;"
                             " 0 = не проверять")
    parser.add_argument("--no-velocity", action="store_true",
                        help="не мерить громкость, оставить константу 100")
    parser.add_argument("--instruments", default="band",
                        help="какие инструменты просить у модели: auto (по умолчанию,"
                             " модель решает сама), band (8 проверенных групп) или список")
    args = parser.parse_args()
    Proxy.target_port = args.target
    Proxy.ping_seconds = args.ping
    Proxy.font_cdn = "" if (args.font_cdn or "").lower() in ("none", "off", "-") else args.font_cdn
    Proxy.measure_velocity = not args.no_velocity
    Proxy.multi_pass, warning = guard_instruments(
        expand_instruments(args.multi_pass), "--multi-pass")
    if warning:
        print(warning, flush=True)
    Proxy.energy_threshold = args.energy_threshold
    Proxy.refine_articulation = not args.no_articulation
    Proxy.instruments, warning = guard_instruments(
        expand_instruments(args.instruments), "--instruments")
    if warning:
        print(warning, flush=True)
    server = ThreadingHTTPServer(("0.0.0.0", args.listen), Proxy)
    server.daemon_threads = True
    print("прокси: 0.0.0.0:%d -> 127.0.0.1:%d" % (args.listen, args.target), flush=True)
    server.serve_forever()


if __name__ == "__main__":
    main()