import argparse
import json
import os
import subprocess
import sys
from pathlib import Path

class LSPClient:
    def __init__(self, server_cmd):
        self.process = subprocess.Popen(
            server_cmd,
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=sys.stderr,
            shell=True
        )
        self.request_id = 0

    def send_message(self, method, params, is_notification=False):
        """Sends a JSON-RPC message."""
        message = {
            "jsonrpc": "2.0",
            "method": method,
            "params": params
        }
        if not is_notification:
            self.request_id += 1
            message["id"] = self.request_id

        body = json.dumps(message)
        header = f"Content-Length: {len(body.encode('utf-8'))}\r\n\r\n"
        
        self.process.stdin.write(header.encode('utf-8'))
        self.process.stdin.write(body.encode('utf-8'))
        self.process.stdin.flush()
        
        return self.request_id

    def read_message(self):
        """Reads a JSON-RPC message."""
        length = 0
        while True:
            line = self.process.stdout.readline().decode('utf-8')
            if not line: raise EOFError("Server closed connection")
            if line.startswith("Content-Length: "):
                length = int(line.split(":")[1].strip())
                break
        
        self.process.stdout.readline() # Skip \r\n
        body = self.process.stdout.read(length).decode('utf-8')
        return json.loads(body)

    def shutdown(self):
        try:
            self.send_message("shutdown", {})
            self.read_message() # Wait for ack
            self.send_message("exit", {}, is_notification=True)
        except:
            pass
        self.process.terminate()

def decode_tokens(data, legend_types, legend_modifiers, source_lines):
    """
    Decodes the flat integer list [deltaLine, deltaStart, len, type, mod]
    into a readable list of tokens.
    """
    tokens = []
    current_line = 0
    current_char = 0
    
    # Process 5 integers at a time
    for i in range(0, len(data), 5):
        delta_line = data[i]
        delta_start = data[i+1]
        length = data[i+2]
        token_type_id = data[i+3]
        token_mod_bitmask = data[i+4]

        # Calculate absolute position
        if delta_line > 0:
            current_line += delta_line
            current_char = delta_start
        else:
            current_char += delta_start

        # Map IDs to names
        type_name = legend_types[token_type_id] if token_type_id < len(legend_types) else str(token_type_id)
        
        # Decode modifiers (bitmask)
        modifiers = []
        for bit in range(len(legend_modifiers)):
            if token_mod_bitmask & (1 << bit):
                modifiers.append(legend_modifiers[bit])

        # Extract the actual text content for verification
        text_content = "EOF/Error"
        if current_line < len(source_lines):
            line_text = source_lines[current_line]
            if current_char + length <= len(line_text) + 1: # +1 for newline tolerance
                text_content = line_text[current_char : current_char + length]

        tokens.append({
            "line": current_line,
            "char": current_char,
            "text": text_content,
            "type": type_name,
            "modifiers": modifiers
        })
    
    return tokens

def main():
    parser = argparse.ArgumentParser(description="Fetch semantic tokens from wasot-lsp.")
    parser.add_argument("filepath", help="Path to the file to analyze")
    args = parser.parse_args()

    file_path = Path(args.filepath).resolve()
    if not file_path.exists():
        print(f"Error: File not found: {file_path}")
        sys.exit(1)

    # Read file content
    with open(file_path, 'r', encoding='utf-8') as f:
        file_content = f.read()
        source_lines = file_content.splitlines()

    client = LSPClient("wasot-lsp")
    
    try:
        # 1. Initialize
        # Request semantic tokens explicitly in capabilities
        init_params = {
            "processId": os.getpid(),
            "rootUri": file_path.parent.as_uri(),
            "capabilities": {
                "textDocument": {
                    "semanticTokens": {
                        "requests": {"full": True},
                        "tokenTypes": [],
                        "tokenModifiers": [],
                        "formats": ["relative"]
                    }
                }
            }
        }
        
        client.send_message("initialize", init_params)
        init_response = client.read_message()
        
        # Extract Legend
        server_caps = init_response['result']['capabilities']
        sem_provider = server_caps.get('semanticTokensProvider', {})
        
        # Handle cases where provider might be a simple ID or full dict
        if not isinstance(sem_provider, dict):
            # Fallback if server doesn't return full provider info in init (rare but possible)
            print("Server enabled semantic tokens but provided no legend in initialize.")
            legend_types = []
            legend_mods = []
        else:
            legend = sem_provider.get('legend', {})
            legend_types = legend.get('tokenTypes', [])
            legend_mods = legend.get('tokenModifiers', [])

        client.send_message("initialized", {}, is_notification=True)

        # 2. Open Document
        client.send_message("textDocument/didOpen", {
            "textDocument": {
                "uri": file_path.as_uri(),
                "languageId": "wat", # Assuming wasot-lsp is for WebAssembly Text
                "version": 1,
                "text": file_content
            }
        }, is_notification=True)

        # 3. Request Tokens
        req_id = client.send_message("textDocument/semanticTokens/full", {
            "textDocument": {"uri": file_path.as_uri()}
        })
        
        # 4. Read Response (Loop to skip potential log messages/diagnostics)
        while True:
            msg = client.read_message()
            if msg.get("id") == req_id:
                token_data = msg.get("result", {}).get("data", [])
                break
        
        # 5. Output
        if not token_data:
            print("No tokens returned.")
        else:
            readable_tokens = decode_tokens(token_data, legend_types, legend_mods, source_lines)
            
            # Print Header
            print(f"{'Line':<6} {'Col':<6} {'Text':<20} {'Type':<20} {'Modifiers'}")
            print("-" * 70)
            
            for t in readable_tokens:
                mod_str = ", ".join(t['modifiers'])
                print(f"{t['line']:<6} {t['char']:<6} {t['text']:<20} {t['type']:<20} {mod_str}")

    except Exception as e:
        print(f"Error: {e}")
    finally:
        client.shutdown()

if __name__ == "__main__":
    main()