#!/usr/bin/env python3
"""Protocol test for the Klammertext language server (klammertext_ls.py).

Spawns the server as an LSP client would (JSON-RPC over stdio) and exercises
every capability against the editor fixtures:

  * initialize / initialized handshake and the advertised capabilities
  * didOpen + publishDiagnostics (clean fixture -> no diagnostics;
    unbalanced text -> the expected errors)
  * textDocument/formatting on every indent fixture == its expected file
  * textDocument/rangeFormatting (a sub-range only)
  * textDocument/documentHighlight + definition (the delimiter matcher)
  * workspace/executeCommand klammertext.alignTable -> workspace/applyEdit,
    applied result == every align fixture's expected file
  * shutdown / exit (exit code 0)

Usage: ls_test.py SHARED_DIR FIXTURE_DIR
Prints one PASS/FAIL line per check; exits nonzero on any failure.
"""

import json
import subprocess
import sys


class Client:
    """A minimal scripted LSP client over a server subprocess."""

    def __init__(self, server_path):
        self.proc = subprocess.Popen(
            [sys.executable, server_path],
            stdin=subprocess.PIPE, stdout=subprocess.PIPE)
        self.next_id = 1
        self.notifications = []      # collected server notifications
        self.server_requests = []    # collected server->client requests

    def send(self, message):
        body = json.dumps(message).encode('utf-8')
        self.proc.stdin.write(b'Content-Length: %d\r\n\r\n' % len(body) + body)
        self.proc.stdin.flush()

    def read_message(self):
        length = None
        while True:
            line = self.proc.stdout.readline()
            if not line:
                return None
            line = line.strip()
            if not line:
                break
            if line.lower().startswith(b'content-length:'):
                length = int(line.split(b':', 1)[1])
        return json.loads(self.proc.stdout.read(length).decode('utf-8'))

    def request(self, method, params):
        """Send a request and pump messages until its response arrives."""
        msg_id = self.next_id
        self.next_id += 1
        self.send({'jsonrpc': '2.0', 'id': msg_id,
                   'method': method, 'params': params})
        while True:
            msg = self.read_message()
            assert msg is not None, 'server closed during ' + method
            if msg.get('id') == msg_id and 'method' not in msg:
                assert 'error' not in msg, 'error response: %r' % msg
                return msg.get('result')
            self._collect(msg)

    def notify(self, method, params):
        self.send({'jsonrpc': '2.0', 'method': method, 'params': params})

    def _collect(self, msg):
        if 'method' in msg and msg.get('id') is not None:
            self.server_requests.append(msg)
            # acknowledge server->client requests (e.g. applyEdit)
            self.send({'jsonrpc': '2.0', 'id': msg['id'],
                       'result': {'applied': True}})
        elif 'method' in msg:
            self.notifications.append(msg)

    def wait_notification(self, method):
        """Pump messages until a notification with METHOD arrives."""
        while True:
            for i, msg in enumerate(self.notifications):
                if msg['method'] == method:
                    return self.notifications.pop(i)['params']
            msg = self.read_message()
            assert msg is not None, 'server closed waiting for ' + method
            self._collect(msg)

    def wait_server_request(self, method):
        while True:
            for i, msg in enumerate(self.server_requests):
                if msg['method'] == method:
                    return self.server_requests.pop(i)['params']
            msg = self.read_message()
            assert msg is not None, 'server closed waiting for ' + method
            self._collect(msg)


PASS = 0
FAIL = 0


def check(name, ok, detail=''):
    global PASS, FAIL
    if ok:
        PASS += 1
        print('PASS %s' % name)
    else:
        FAIL += 1
        print('FAIL %s %s' % (name, detail))


def apply_edits(ls, text, edits):
    """Apply LSP TextEdits to TEXT (offsets via the server's own helpers)."""
    resolved = [(ls.pos_to_offset(text, e['range']['start']),
                 ls.pos_to_offset(text, e['range']['end']),
                 e['newText']) for e in edits]
    for a, b, new in sorted(resolved, reverse=True):
        text = text[:a] + new + text[b:]
    return text


def main():
    shared_dir, fixture_dir = sys.argv[1], sys.argv[2]
    sys.path.insert(0, shared_dir)
    import klammertext_ls as ls

    client = Client(shared_dir + '/klammertext_ls.py')
    uri = 'file:///ls_test.kt'

    # -- handshake --
    result = client.request('initialize', {'capabilities': {}})
    caps = result.get('capabilities', {})
    check('initialize capabilities',
          caps.get('documentFormattingProvider') is True
          and caps.get('documentHighlightProvider') is True
          and 'klammertext.alignTable'
              in caps.get('executeCommandProvider', {}).get('commands', []))
    client.notify('initialized', {})

    # -- diagnostics: clean fixture, then unbalanced text --
    text = open(fixture_dir + '/indent_list_expected.kt').read()
    client.notify('textDocument/didOpen',
                  {'textDocument': {'uri': uri, 'languageId': 'klammertext',
                                    'version': 1, 'text': text}})
    diags = client.wait_notification('textDocument/publishDiagnostics')
    check('clean fixture: no diagnostics', diags['diagnostics'] == [],
          repr(diags['diagnostics']))

    bad = '@i abc\n@ol x ul@\n'
    client.notify('textDocument/didChange',
                  {'textDocument': {'uri': uri, 'version': 2},
                   'contentChanges': [{'text': bad}]})
    diags = client.wait_notification('textDocument/publishDiagnostics')
    msgs = [d['message'] for d in diags['diagnostics']]
    check('unbalanced text: diagnostics',
          any('@i' in m and 'never closed' in m for m in msgs)
          and any('ul@' in m for m in msgs), repr(msgs))

    # -- formatting == every indent fixture's expected file --
    indent_fixtures = ['indent_list', 'indent_document', 'indent_table',
                       'indent_untouched', 'indent_defs', 'indent_escapes',
                       'indent_named_close']
    for f in indent_fixtures:
        src = open(fixture_dir + '/%s.kt' % f).read()
        exp = open(fixture_dir + '/%s_expected.kt' % f).read()
        client.notify('textDocument/didChange',
                      {'textDocument': {'uri': uri, 'version': 3},
                       'contentChanges': [{'text': src}]})
        client.wait_notification('textDocument/publishDiagnostics')
        edits = client.request('textDocument/formatting',
                               {'textDocument': {'uri': uri},
                                'options': {'tabSize': 2,
                                            'insertSpaces': True}})
        check('formatting %s' % f, apply_edits(ls, src, edits or []) == exp)

    # -- rangeFormatting: only the requested lines change --
    src = '@ol\nzero\n    one\n@\n'
    client.notify('textDocument/didChange',
                  {'textDocument': {'uri': uri, 'version': 4},
                   'contentChanges': [{'text': src}]})
    client.wait_notification('textDocument/publishDiagnostics')
    edits = client.request('textDocument/rangeFormatting',
                           {'textDocument': {'uri': uri},
                            'range': {'start': {'line': 2, 'character': 0},
                                      'end': {'line': 2, 'character': 0}},
                            'options': {}})
    check('rangeFormatting touches only its lines',
          apply_edits(ls, src, edits or []) == '@ol\nzero\n  one\n@\n')

    # -- matcher: documentHighlight + definition --
    src = '@i abc @\n'
    client.notify('textDocument/didChange',
                  {'textDocument': {'uri': uri, 'version': 5},
                   'contentChanges': [{'text': src}]})
    client.wait_notification('textDocument/publishDiagnostics')
    hl = client.request('textDocument/documentHighlight',
                        {'textDocument': {'uri': uri},
                         'position': {'line': 0, 'character': 0}})
    check('documentHighlight pair', hl is not None and len(hl) == 2, repr(hl))
    defn = client.request('textDocument/definition',
                          {'textDocument': {'uri': uri},
                           'position': {'line': 0, 'character': 0}})
    check('definition = the matching close',
          defn is not None and defn['range']['start']['character'] == 7,
          repr(defn))
    off = client.request('textDocument/documentHighlight',
                         {'textDocument': {'uri': uri},
                          'position': {'line': 0, 'character': 4}})
    check('documentHighlight off-delimiter is null', off is None, repr(off))

    # -- the custom matchInfo request (drives the VS Code decorations) --
    hl2 = client.request('textDocument/documentHighlight',
                         {'textDocument': {'uri': uri},
                          'position': {'line': 0, 'character': 7}})
    check('documentHighlight from the bare close',
          hl2 is not None and len(hl2) == 2, repr(hl2))
    mi = client.request('klammertext/matchInfo',
                        {'textDocument': {'uri': uri},
                         'position': {'line': 0, 'character': 7}})
    check('matchInfo from the bare close',
          mi is not None and not mi['mismatch']
          and mi['matchToken']['start']['character'] == 0, repr(mi))
    client.notify('textDocument/didChange',
                  {'textDocument': {'uri': uri, 'version': 6},
                   'contentChanges': [{'text': '@ol x ul@\n'}]})
    client.wait_notification('textDocument/publishDiagnostics')
    mi = client.request('klammertext/matchInfo',
                        {'textDocument': {'uri': uri},
                         'position': {'line': 0, 'character': 0}})
    check('matchInfo reports a mismatch',
          mi is not None and mi['mismatch'] and 'ul@' in (mi['message'] or ''),
          repr(mi))
    mi = client.request('klammertext/matchInfo',
                        {'textDocument': {'uri': uri},
                         'position': {'line': 0, 'character': 4}})
    check('matchInfo off-delimiter is null', mi is None, repr(mi))

    # -- alignTable via executeCommand -> applyEdit, on every align fixture --
    align_fixtures = ['align_mixed', 'align_empty_cells', 'align_boundary',
                      'align_colspan', 'align_escapes']
    for f in align_fixtures:
        src = open(fixture_dir + '/%s.kt' % f).read()
        exp = open(fixture_dir + '/%s_expected.kt' % f).read()
        client.notify('textDocument/didChange',
                      {'textDocument': {'uri': uri, 'version': 6},
                       'contentChanges': [{'text': src}]})
        client.wait_notification('textDocument/publishDiagnostics')
        pos = ls.offset_to_pos(src, src.index('|'))
        client.request('workspace/executeCommand',
                       {'command': 'klammertext.alignTable',
                        'arguments': [{'uri': uri, 'position': pos}]})
        if src == exp:               # already-aligned fixture: no applyEdit
            continue
        params = client.wait_server_request('workspace/applyEdit')
        edits = params['edit']['changes'][uri]
        check('alignTable %s' % f, apply_edits(ls, src, edits) == exp)

    # align_too_wide must produce NO applyEdit (limit refusal) — verify via
    # a message arriving with no pending applyEdit request.
    src = open(fixture_dir + '/align_too_wide.kt').read()
    client.notify('textDocument/didChange',
                  {'textDocument': {'uri': uri, 'version': 7},
                   'contentChanges': [{'text': src}]})
    client.wait_notification('textDocument/publishDiagnostics')
    client.notifications.clear()     # drop showMessages from the loop above
    pos = ls.offset_to_pos(src, src.index('|'))
    client.request('workspace/executeCommand',
                   {'command': 'klammertext.alignTable',
                    'arguments': [{'uri': uri, 'position': pos}]})
    msg = client.wait_notification('window/showMessage')
    check('alignTable respects ROW_MAX (refusal message, no edit)',
          'not aligning' in msg['message'] and not client.server_requests,
          repr(msg))

    # -- shutdown --
    client.request('shutdown', None)
    client.notify('exit', None)
    code = client.proc.wait(timeout=10)
    check('clean exit', code == 0, 'exit code %d' % code)

    print('ls_test: %d passed, %d failed' % (PASS, FAIL))
    sys.exit(1 if FAIL else 0)


if __name__ == '__main__':
    main()
