#!/usr/bin/env python3

import ast
import argparse
import json
from pathlib import Path

DEFAULT_REPO = "/home/uha4/tmp/tryton"

DEFAULT_MODULES = [
    "ir",
    "res",
    "company",
    "country",
    "currency",
    "party",
    "product",
    "production_split",
    "stock_lot",
    "stock_lot_sled",
    "production",
    "stock_product_location",
    "stock_split",
    "product_measurements",
]

DEFAULT_OUTPUT = "data"

class RegisterExtractor(ast.NodeVisitor):

    def __init__(self, filename):
        self.filename = filename
        self.results = []
        self.model_names = {}

    def visit_Module(self, node):
        for item in node.body:
            if isinstance(item, ast.ClassDef):
                model_name = self.get_model_name(item)
                if model_name:
                    self.model_names[item.name] = model_name

        self.generic_visit(node)

    def visit_ClassDef(self, node):
        for item in node.body:
            if isinstance(item, ast.FunctionDef) and item.name == "__register__":
                self.extract_register(node.name, item)

        self.generic_visit(node)

    def extract_register(self, cls, func):

        operations = []
        self.table_vars = {}
        self.current_class = cls

        model_name = self.model_names.get(cls)
        if model_name:
            self.table_vars["cls"] = model_name

        for stmt in func.body:
            self.collect_table_assignment(stmt)

        self.walk_statements(func.body, operations, "before_super", None)

        self.results.append({
            "file": self.filename,
            "class": cls,
            "method": "__register__",
            "line": func.lineno,
            "operations": operations,
        })


    def get_model_name(self, node):
        for item in node.body:
            if not isinstance(item, ast.Assign):
                continue

            for target in item.targets:
                if (
                    isinstance(target, ast.Name)
                    and target.id == "__name__"
                ):
                    try:
                        return ast.literal_eval(item.value)
                    except Exception:
                        return None

        return None

    def collect_table_assignment(self, stmt):
        for node in ast.walk(stmt):
            if not isinstance(node, ast.Assign):
                continue

            if len(node.targets) != 1:
                continue

            target = node.targets[0]
            if not isinstance(target, ast.Name):
                continue

            table = self.resolve_table_call(node.value)
            if table:
                self.table_vars[target.id] = table
            elif isinstance(node.value, ast.Constant):
                self.table_vars[target.id] = node.value.value

    def resolve_table_call(self, node):
        if not isinstance(node, ast.Call):
            return None

        # ModelClass = pool.get('ir.model.data')
        if (
            isinstance(node.func, ast.Attribute)
            and node.func.attr == "get"
            and isinstance(node.func.value, ast.Name)
            and node.func.value.id == "pool"
            and len(node.args) == 1
            and isinstance(node.args[0], ast.Constant)
            and isinstance(node.args[0].value, str)
        ):
            return node.args[0].value


        if not (
            isinstance(node.func, ast.Attribute)
            and node.func.attr in ("__table__", "__table_handler__")
        ):
            return None

        receiver = node.func.value
        # ModelClass = pool.get('ir.model.data')
        if (
            isinstance(receiver, ast.Name)
            and receiver.id == "pool"
        ):
            return None

        if (
            isinstance(node.func, ast.Attribute)
            and node.func.attr == "get"
            and isinstance(node.func.value, ast.Name)
            and node.func.value.id == "pool"
            and len(node.args) == 1
            and isinstance(node.args[0], ast.Constant)
            and isinstance(node.args[0].value, str)
        ):
            return node.args[0].value
        if isinstance(receiver, ast.Name):
            if receiver.id == "cls":
                model_name = self.model_names.get(self.current_class)
            else:
                model_name = (
                    self.model_names.get(receiver.id)
                    or self.table_vars.get(receiver.id)
                )
                
            if model_name:
                return model_name.replace(".", "_")

        return None

    def resolve_table_object(self, node):
        if isinstance(node, ast.Name):
            return self.table_vars.get(node.id)
    
        if (
            isinstance(node, ast.Attribute)
            and node.attr == "_table"
            and isinstance(node.value, ast.Name)
            and node.value.id == "cls"
        ):
            model_name = self.model_names.get(self.current_class)
            if model_name:
                return model_name.replace(".", "_")
    
        return self.resolve_table_call(node)

    def find_tables(self, node):
        tables = []

        for child in ast.walk(node):
            if not isinstance(child, ast.Call):
                continue

            if not (
                isinstance(child.func, ast.Attribute)
                and child.func.attr in {
                    "column_rename",
                    "drop_column",
                    "add_column",
                    "drop_constraint",
                    "add_constraint",
                    "not_null_action",
                    "index_action",
                    "alter_size",
                    "update",
                    "insert",
                    "delete",
                    "select",
                    "table_exist",
                }
            ):
                continue
            
            # backend.TableHandler.table_exist(cls._table)
            if (
                isinstance(child.func.value, ast.Attribute)
                and child.func.value.attr == "TableHandler"
                and isinstance(child.func.value.value, ast.Name)
                and child.func.value.value.id == "backend"
                and child.func.attr == "table_exist"
                and child.args
            ):
                table = self.resolve_table_object(child.args[0])
                if table and table not in tables:
                    tables.append(table)
    
                continue
            table = self.resolve_table_object(child.func.value)
            if table and table not in tables:
                tables.append(table)

        return tables

    def walk_statements(self, statements, operations, phase, condition):

        for stmt in statements:

            if self.is_super_register(stmt):
                phase = "after_super"
                continue

            if isinstance(stmt, ast.If):
                cond = ast.unparse(stmt.test)
                self.walk_statements(stmt.body, operations, phase, cond)
                self.walk_statements(stmt.orelse, operations, phase, condition)
                continue

            value = getattr(stmt, "value", None)
            if value is not None:
                self.walk_expression(value, operations, phase, condition)

    def walk_expression(self, node, operations, phase, condition):

        if node is None:
            return

        if isinstance(node, ast.Call):
            op = self.extract_call(node)
            if op:
                op["phase"] = phase
                if condition:
                    op["condition"] = condition
                operations.append(op)
                
                if not op['tables']:
                    print(op)

        for child in ast.iter_child_nodes(node):
            self.walk_expression(child, operations, phase, condition)

    def is_super_register(self, stmt):

        if not isinstance(stmt, ast.Expr):
            return False

        call = stmt.value

        if not isinstance(call, ast.Call):
            return False

        func = call.func

        if not isinstance(func, ast.Attribute):
            return False

        if func.attr != "__register__":
            return False

        if not isinstance(func.value, ast.Call):
            return False

        if not isinstance(func.value.func, ast.Name):
            return False

        return func.value.func.id == "super"

    def extract_call(self, call):
    
        if not isinstance(call.func, ast.Attribute):
            return None
    
        interesting = {
            "column_rename",
            "drop_column",
            "add_column",
            "drop_constraint",
            "add_constraint",
            "not_null_action",
            "index_action",
            "alter_size",
            "execute",
            "executemany",
            "update",
            "insert",
            "delete",
            "select",
        }
    
        operation = call.func.attr
    
        is_table_handler = (
            isinstance(call.func.value, ast.Attribute)
            and call.func.value.attr == "TableHandler"
            and isinstance(call.func.value.value, ast.Name)
            and call.func.value.value.id == "backend"
        )

        if operation not in interesting and not is_table_handler:
            return None

        tables = self.find_tables(call)
        
        if (
            isinstance(call.func, ast.Attribute)
            and call.func.attr == "table_rename"
            and isinstance(call.func.value, ast.Attribute)
            and isinstance(call.func.value.value, ast.Name)
            and call.func.value.value.id == "backend"
            and isinstance(call.args, list)
        ):
            for arg in call.args[:2]:
                if isinstance(arg, ast.Constant) and isinstance(arg.value, str):
                    table = arg.value
                elif isinstance(arg, ast.Name):
                    table = self.table_vars.get(arg.id)
                elif (
                    isinstance(arg, ast.Attribute)
                    and isinstance(arg.value, ast.Name)
                    and arg.value.id == "cls"
                    and arg.attr == "_table"
                ):
                    table = self.model_names.get(self.current_class)
                    if table:
                        table = table.replace(".", "_")
                else:
                    table = None

                if table and table not in tables:
                    tables.append(table)    
                    
        # Wer wird aufgerufen?
        if isinstance(call.func.value, ast.Name):
            obj = call.func.value.id
        else:
            obj = ast.unparse(call.func.value)
    
        args = []
    
        for arg in call.args:
            try:
                args.append(ast.unparse(arg))
            except Exception:
                args.append("<unknown>")
    
        kwargs = {}
    
        for kw in call.keywords:
            try:
                kwargs[kw.arg] = ast.unparse(kw.value)
            except Exception:
                kwargs[kw.arg] = "<unknown>"
                
        return {
            "object": obj,
            "operation": operation,
            "tables": tables,
            "line": call.lineno,
            "args": args,
            "kwargs": kwargs,
            "code": ast.unparse(call),
        }

def analyse_file(path):

    try:
        tree = ast.parse(path.read_text(encoding="utf-8"))
    except Exception as e:
        print(f"Cannot parse {path}: {e}")
        return []

    visitor = RegisterExtractor(str(path))
    visitor.visit(tree)

    results = [r for r in visitor.results if r["operations"]]
    
    return results

def detect_version(repo: Path) -> str:
    changelog = repo / "trytond" / "CHANGELOG"

    if not changelog.exists():
        return "unknown"

    with changelog.open("r", encoding="utf-8") as f:
        for line in f:
            if line.startswith("Version "):
                return line.split()[1].rsplit('.', 1)[0]

    return "unknown"

def find_module(repo: Path, module: str) -> Path | None:

    candidates = [
        repo / "modules" / module,
        repo / "trytond" / "trytond" / module,
    ]

    for c in candidates:
        if c.exists():
            return c

    return None

def main():

    parser = argparse.ArgumentParser()

    parser.add_argument(
        "--repo",
        help="Path to Tryton repository",
    )

    parser.add_argument(
        "--modules",
        nargs="+",
        help="Module names",
    )

    parser.add_argument(
        "--output",
        help="Output directory",
    )

    args = parser.parse_args()

    repo = Path(args.repo or DEFAULT_REPO)

    modules = args.modules or DEFAULT_MODULES

    output = Path(args.output or DEFAULT_OUTPUT)

    version = detect_version(repo)

    print(f"Repository : {repo}")
    print(f"Version    : {version}")
    print(f"Modules    : {', '.join(modules)}")
    print()

    for module in modules:

        module_path = find_module(repo, module)

        if not module_path:
            print(f"Module not found: {module}")
            continue

        results = []

        for py in sorted(module_path.rglob("*.py")):
            results.extend(analyse_file(py))

        outdir = output / version
        outdir.mkdir(parents=True, exist_ok=True)

        outfile = outdir / f"{module}.json"

        with outfile.open("w", encoding="utf-8") as f:
            json.dump(
                {
                    "version": version,
                    "module": module,
                    "files": results,
                },
                f,
                indent=2,
            )

        print(f"{module:25} -> {outfile}")


if __name__ == "__main__":
    main()