#!/usr/bin/env python3
"""
Minimal RPC code generator PoC.

Parses C++ sources with libclang, finds classes and methods annotated with
ANNOTATE("export") (or [[annotate("export")]]), and generates proxy + skeleton
for methods that use only `int` arguments and return type `int`.

Usage:
    tools/generate_rpc.py --out-dir build/generated --compile-commands build/compile_commands.json src/MyService.h src/other.h
"""
import argparse
import json
import os
import sys
from dataclasses import dataclass, field
from typing import List

# Try imports
try:
    from clang.cindex import Index, CursorKind, Config
except Exception as e:
    print("ERROR: clang.cindex import failed:", e)
    print("Make sure libclang python bindings are installed and libclang is reachable.")
    sys.exit(1)

try:
    from jinja2 import Environment, FileSystemLoader
except Exception as e:
    print("ERROR: jinja2 import failed:", e)
    print("Install with: pip3 install jinja2")
    sys.exit(1)

# ========== IR dataclasses ==========
@dataclass
class Arg:
    name: str
    type: str

@dataclass
class Method:
    name: str
    args: List[Arg] = field(default_factory=list)
    return_type: str = "void"
    signature: str = ""   # unique signature string used for dispatch (name:types)

@dataclass
class Class:
    name: str
    methods: List[Method] = field(default_factory=list)

# ========== helpers for compile flags ==========
def load_compile_flags(compile_commands_path, src_path):
    if not compile_commands_path or not os.path.exists(compile_commands_path):
        return []
    with open(compile_commands_path) as f:
        compile_commands = json.load(f)
    # find matching entry by file (exact or basename)
    abs_src = os.path.abspath(src_path)
    for entry in compile_commands:
        entry_file = os.path.abspath(entry.get('file') or "")
        if entry_file == abs_src or os.path.basename(entry_file) == os.path.basename(abs_src):
            # get arguments
            if 'arguments' in entry and isinstance(entry['arguments'], list):
                cmd = entry['arguments']
            else:
                cmd = (entry.get('command') or "").split()
            # remove compiler and -c <file>
            # keep flags like -I, -D
            filtered = []
            skip_next = False
            for i, tok in enumerate(cmd):
                if skip_next:
                    skip_next = False
                    continue
                # skip the compiler binary itself
                if (
                    tok.endswith("g++")
                    or tok.endswith("clang++")
                    or tok.endswith("clang")
                    or tok.endswith("cc")
                    or tok.endswith("c++")
                ):
                    continue
                if tok == "-c":
                    skip_next = True
                    continue
                # drop output file flag: -o <file>
                if tok == "-o":
                    skip_next = True
                    continue
                filtered.append(tok)
            return filtered
    return []

# ========== AST utilities ==========
def has_annotation(cursor, annotation_text="export"):
    # Check children for ANNOTATE_ATTR or spelling containing annotation_text
    for c in cursor.get_children():
        if c.kind == CursorKind.ANNOTATE_ATTR and annotation_text in (c.spelling or ""):
            return True
    # also check attributes spelling (fallback)
    # cursor.get_tokens() might include annotate attribute for new attribute syntax but libclang may not expose
    return False

def type_is_int(t):
    if not t:
        return False
    # check canonical type spelling - this may vary; keep simple
    s = t.spelling
    # accept "int" and "signed int"
    return s in ("int", "signed int", "int32_t", "long int")  # minimal

# ========== parse file ==========
def parse_file(index, filepath, args):
    tu = index.parse(filepath, args=args)
    classes = []

    for cursor in tu.cursor.get_children():
        # top-level class/struct declarations
        if cursor.kind == CursorKind.CLASS_DECL or cursor.kind == CursorKind.STRUCT_DECL:
            # only records with a name
            if not cursor.spelling:
                continue
            if has_annotation(cursor, "export"):
                cls = Class(name=cursor.spelling)
                # iterate its children for methods
                for c in cursor.get_children():
                    if c.kind == CursorKind.CXX_METHOD or c.kind == CursorKind.FUNCTION_DECL:
                        if not has_annotation(c, "export"):
                            continue
                        # build method IR
                        # only support non-template, non-variadic, simple methods
                        m = Method(name=c.spelling)
                        # return type
                        try:
                            m.return_type = c.result_type.spelling
                        except Exception:
                            m.return_type = "void"
                        # args
                        ok = True
                        for arg in c.get_arguments():
                            t = arg.type
                            if not type_is_int(t):
                                ok = False
                                break
                            m.args.append(Arg(name=arg.spelling or "arg", type=t.spelling))
                        if not ok:
                            print(f"Skipping method {c.spelling} of {cursor.spelling}: unsupported arg types")
                            continue
                        # check return type is int
                        if not type_is_int(c.result_type):
                            print(f"Skipping method {c.spelling} of {cursor.spelling}: unsupported return type {c.result_type.spelling}")
                            continue
                        # signature: name:comma-separated types
                        sig_types = ",".join(a.type for a in m.args)
                        m.signature = f"{m.name}:{sig_types}"
                        cls.methods.append(m)
                if cls.methods:
                    classes.append(cls)
    return classes

# ========== templating ==========
def render_templates(classes, out_dir, templates_dir):
    env = Environment(
        loader=FileSystemLoader(templates_dir),
        autoescape=False,
        trim_blocks=True,
        lstrip_blocks=True,
    )

    proxy_h = env.get_template("proxy.h.j2")
    proxy_cpp = env.get_template("proxy.cpp.j2")
    skeleton_h = env.get_template("skeleton.h.j2")
    skeleton_cpp = env.get_template("skeleton.cpp.j2")

    for cls in classes:
        name = cls.name

        with open(f"{out_dir}/{name}.proxy.h", "w") as f:
            f.write(proxy_h.render(cls=cls))

        with open(f"{out_dir}/{name}.proxy.cpp", "w") as f:
            f.write(proxy_cpp.render(cls=cls))

        with open(f"{out_dir}/{name}.skeleton.h", "w") as f:
            f.write(skeleton_h.render(cls=cls))

        with open(f"{out_dir}/{name}.skeleton.cpp", "w") as f:
            f.write(skeleton_cpp.render(cls=cls))

# ========== main ==========
def main():
    p = argparse.ArgumentParser()
    p.add_argument("--out-dir", "-o", required=True)
    p.add_argument("--compile-commands", "-c", default=None)
    p.add_argument("--header", required=True, help="C++ header file to scan for exported RPC classes")
    p.add_argument("--source", required=True, help="C++ source file to use for resolving compile flags (from compile_commands.json)")
    p.add_argument(
        "--out-base",
        required=True,
        help="Base name for generated files: <out-base>.proxy.[h|cpp], <out-base>.skeleton.[h|cpp]",
    )
    p.add_argument("--templates", "-t", default=os.path.join(os.path.dirname(__file__), "templates"))
    args = p.parse_args()

    out_dir = args.out_dir
    os.makedirs(out_dir, exist_ok=True)

    # configure libclang: optionally allow user to set LIBCLANG_PATH env var
    # If needed, user can set: export LIBCLANG_PATH=/usr/lib/llvm-10/lib
    libclang_path = os.environ.get("LIBCLANG_PATH")
    if libclang_path:
        try:
            Config.set_library_path(libclang_path)
        except Exception as e:
            print("WARNING: cannot set libclang path:", e)

    # determine compile flags from the source file (which is what appears in compile_commands.json)
    if not os.path.exists(args.header):
        print("ERROR: header file not found:", args.header)
        return 1
    if not os.path.exists(args.source):
        print("ERROR: source file not found:", args.source)
        return 1

    compile_args = load_compile_flags(args.compile_commands, args.source)
    if not any(a.startswith("-x") for a in compile_args):
        compile_args = ["-x", "c++", "-std=c++17"] + compile_args

    print("Parsing", args.header, "with args:", compile_args)

    index = Index.create()
    classes = parse_file(index, args.header, compile_args)

    if not classes:
        print("No exported classes/methods found in header. Nothing to generate.")
        return 0

    # For this PoC we expect a single exported service per header.
    # If there are multiple, refuse to guess which one should define the filenames.
    if len(classes) > 1:
        print(
            "ERROR: multiple exported classes found in header; "
            "current generator expects exactly one when using --out-base."
        )
        for c in classes:
            print(" -", c.name)
        return 1

    # render templates using the single discovered class but the user‑provided base name
    cls = classes[0]

    env = Environment(
        loader=FileSystemLoader(args.templates),
        autoescape=False,
        trim_blocks=True,
        lstrip_blocks=True,
    )

    proxy_h = env.get_template("proxy.h.j2")
    proxy_cpp = env.get_template("proxy.cpp.j2")
    skeleton_h = env.get_template("skeleton.h.j2")
    skeleton_cpp = env.get_template("skeleton.cpp.j2")

    base = args.out_base

    with open(f"{out_dir}/{base}.proxy.h", "w") as f:
        f.write(proxy_h.render(cls=cls))

    with open(f"{out_dir}/{base}.proxy.cpp", "w") as f:
        f.write(proxy_cpp.render(cls=cls))

    with open(f"{out_dir}/{base}.skeleton.h", "w") as f:
        f.write(skeleton_h.render(cls=cls))

    with open(f"{out_dir}/{base}.skeleton.cpp", "w") as f:
        f.write(skeleton_cpp.render(cls=cls))

    print("Generated files for class", cls.name, "into base", base)
    return 0

if __name__ == "__main__":
    sys.exit(main())