#!/usr/bin/env python3
"""
Python AST Worker - JSON-RPC service for parsing Python code

This worker reads JSON-RPC requests from stdin and writes responses to stdout.
It uses Python's built-in ast module to parse Python code and extract:
- Function definitions (sync and async)
- Class definitions with methods
- Import statements
- Function calls within code

No external dependencies required - uses only Python standard library.
"""
import sys
import io
import json
import ast
from typing import List, Dict, Any

# Fix encoding on Windows: Node.js sends UTF-8 JSON over stdin pipes, but
# Python on Windows defaults to the system ANSI code page (e.g. cp1252).
# This causes multi-byte UTF-8 sequences (like smart quotes \xe2\x80\x9d)
# to be misread, producing surrogate characters (\udc9d) that later crash
# json.dumps with "surrogates not allowed".
if sys.platform == 'win32':
    sys.stdin = io.TextIOWrapper(sys.stdin.buffer, encoding='utf-8', errors='replace')
    sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8', errors='replace', line_buffering=True)


def is_exported(node: ast.AST) -> bool:
    """Check if a function or class is exported (not private).

    In Python, names starting with '_' are conventionally private.
    """
    if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
        return not node.name.startswith('_')
    return False


def get_signature(node: ast.FunctionDef | ast.AsyncFunctionDef) -> str:
    """Extract function signature from AST node."""
    args_list = []

    for arg in node.args.args:
        arg_str = arg.arg
        if arg.annotation:
            arg_str += f': {ast.unparse(arg.annotation)}'
        args_list.append(arg_str)

    return_annotation = ''
    if node.returns:
        return_annotation = f' -> {ast.unparse(node.returns)}'

    return f"{node.name}({', '.join(args_list)}){return_annotation}"


def extract_imports(tree: ast.AST) -> List[Dict[str, Any]]:
    """Extract import statements from AST."""
    imports = []

    for node in ast.walk(tree):
        if isinstance(node, ast.Import):
            for alias in node.names:
                imports.append({
                    'source': alias.name,
                    'imported': alias.asname if alias.asname else alias.name
                })
        elif isinstance(node, ast.ImportFrom):
            module = node.module if node.module else ''
            for alias in node.names:
                imports.append({
                    'source': module,
                    'imported': alias.name,
                    'alias': alias.asname if alias.asname else None
                })

    return imports


def extract_calls(node: ast.AST) -> List[str]:
    """Extract function calls from a function/method body."""
    calls = []

    for child in ast.walk(node):
        if isinstance(child, ast.Call):
            if isinstance(child.func, ast.Name):
                calls.append(child.func.id)
            elif isinstance(child.func, ast.Attribute):
                calls.append(child.func.attr)

    return calls


def parse_python_ast(code: str, file_path: str) -> Dict[str, Any]:
    """Parse Python code and return CodeNode array.

    Args:
        code: Python source code to parse
        file_path: Path to the file (for error messages)

    Returns:
        Dictionary with 'nodes' and 'imports' arrays

    Raises:
        Exception: If parsing fails due to syntax errors
    """
    try:
        tree = ast.parse(code)
        nodes = []

        for node in tree.body:
            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
                nodes.append({
                    'type': 'function',
                    'name': node.name,
                    'exported': is_exported(node),
                    'startLine': node.lineno,
                    'endLine': node.end_lineno if node.end_lineno else node.lineno,
                    'async': isinstance(node, ast.AsyncFunctionDef),
                    'signature': get_signature(node),
                    'calls': extract_calls(node)
                })

            elif isinstance(node, ast.ClassDef):
                methods = []
                for item in node.body:
                    if isinstance(item, (ast.FunctionDef, ast.AsyncFunctionDef)):
                        methods.append({
                            'name': item.name,
                            'async': isinstance(item, ast.AsyncFunctionDef),
                            'signature': get_signature(item),
                            'startLine': item.lineno,
                            'endLine': item.end_lineno if item.end_lineno else item.lineno,
                            'calls': extract_calls(item)
                        })

                nodes.append({
                    'type': 'class',
                    'name': node.name,
                    'exported': is_exported(node),
                    'startLine': node.lineno,
                    'endLine': node.end_lineno if node.end_lineno else node.lineno,
                    'methods': methods
                })

        imports = extract_imports(tree)

        return {
            'nodes': nodes,
            'imports': imports
        }

    except SyntaxError as e:
        raise Exception(f"Python syntax error at line {e.lineno}: {e.msg}")
    except Exception as e:
        raise Exception(f"Failed to parse Python AST: {str(e)}")


def process_parse_python(request: Dict[str, Any]) -> None:
    """Process a parse_python JSON-RPC request."""
    try:
        params = request.get('params', {})
        code = params.get('code')
        file_path = params.get('filePath', '<unknown>')

        if code is None:
            raise ValueError('code parameter is required')

        result = parse_python_ast(code, file_path)
        response = {
            'jsonrpc': '2.0',
            'id': request.get('id'),
            'result': result
        }
        print(json.dumps(response), flush=True)
    except Exception as e:
        error_response = {
            'jsonrpc': '2.0',
            'id': request.get('id'),
            'error': {'code': -1, 'message': str(e)}
        }
        print(json.dumps(error_response), flush=True)


def main():
    """Main loop processing stdin requests."""
    for line in sys.stdin:
        try:
            request = json.loads(line.strip())
            method = request.get('method')

            if method == 'parse_python':
                process_parse_python(request)
            else:
                error_response = {
                    'jsonrpc': '2.0',
                    'id': request.get('id'),
                    'error': {'code': -32601, 'message': f'Unknown method: {method}'}
                }
                print(json.dumps(error_response), flush=True)

        except json.JSONDecodeError as e:
            error_response = {
                'jsonrpc': '2.0',
                'id': None,
                'error': {'code': -32700, 'message': f'Parse error: {str(e)}'}
            }
            print(json.dumps(error_response), flush=True)
        except Exception as e:
            error_response = {
                'jsonrpc': '2.0',
                'id': request.get('id') if 'request' in dir() and isinstance(request, dict) else None,
                'error': {'code': -1, 'message': str(e)}
            }
            print(json.dumps(error_response), flush=True)


if __name__ == '__main__':
    main()