#!/usr/bin/env python3
"""
AST-based transformer to extract App._render_xxx methods into module-level functions.

Transformation:
 class App:
  def _render_xxx(self, ...) -> ReturnType:
   self.foo = bar
   self.baz.qux()

 Becomes:

 def render_xxx(app: App, ...) -> ReturnType:
  app.foo = bar
  app.baz.qux()

 class App:
  def _render_xxx(self, ...) -> ReturnType:
   render_xxx(self, ...)
"""

import ast
import sys
import argparse
from pathlib import Path
from typing import Optional


class RenderMethodTransformer(ast.NodeTransformer):
 def __init__(self, class_name: str = "App", method_prefix: str = "_render"):
  self.class_name = class_name
  self.method_prefix = method_prefix
  self.transformed_methods: list[str] = []
  self.in_target_class = False
  self.current_method: Optional[str] = None

 def visit_ClassDef(self, node: ast.ClassDef) -> ast.ClassDef:
  if node.name == self.class_name:
   self.in_target_class = True
   self.current_method = None
   node = self.generic_visit(node)
   self.in_target_class = False
   return node
  return node

 def visit_FunctionDef(self, node: ast.FunctionDef) -> ast.FunctionDef:
  if not self.in_target_class:
   return node

  if not node.name.startswith(self.method_prefix):
   return node

  if node.name == "__init__" or node.name.startswith("__"):
   return node

  self.current_method = node.name
  new_name = node.name[1:]  # Remove underscore prefix: _render_xxx -> render_xxx

  self.transformed_methods.append(node.name)

  new_func = ast.FunctionDef(
   name=new_name,
   args=ast.arguments(
    posonlyargs=[],
    args=[ast.arg(arg="app", annotation=ast.Name(id=self.class_name, ctx=ast.Load()))] + node.args.args,
    kwonlyargs=node.args.kwonlyargs,
    kw_defaults=node.args.kw_defaults,
    defaults=node.args.defaults,
   ),
   body=[self._transform_body(body_item) for body_item in node.body],
   decorator_list=node.decorator_list,
   returns=node.returns,
   type_comment=node.type_comment,
  )

  delegation_method = ast.FunctionDef(
   name=node.name,
   args=node.args,
   body=[ast.Expr(value=ast.Call(
    func=ast.Name(id=new_name, ctx=ast.Load()),
    args=[ast.Name(id="self", ctx=ast.Load())] + [self._arg_to_expr(arg) for arg in node.args.args],
    keywords=[ast.keyword(arg=arg.arg, value=ast.Name(id=arg.arg, ctx=ast.Load())) for arg in node.args.kwonlyargs]
   ))],
   decorator_list=node.decorator_list,
   returns=node.returns,
   type_comment=node.type_comment,
  )

  return ast.Module(body=[new_func, delegation_method], type_ignores=[])

 def _arg_to_expr(self, arg: ast.arg) -> ast.expr:
  return ast.Name(id=arg.arg, ctx=ast.Load())

 def _transform_body(self, node: ast.AST) -> ast.AST:
  return ast.NodeTransformer.generic_visit(self, node)

 def visit_Name(self, node: ast.Name) -> ast.Name:
  if isinstance(node.ctx, ast.Load) and node.id == "self":
   return ast.Name(id="app", ctx=node.ctx)
  return node

 def visit_Attribute(self, node: ast.Attribute) -> ast.Attribute:
  if isinstance(node.value, ast.Name) and node.value.id == "self":
   new_value = ast.Name(id="app", ctx=node.value.ctx)
   return ast.Attribute(value=new_value, attr=node.attr, ctx=node.ctx)
  return self.generic_visit(node)


class AppRenderExtractor(ast.NodeVisitor):
 def __init__(self, class_name: str = "App", method_prefix: str = "_render"):
  self.class_name = class_name
  self.method_prefix = method_prefix
  self.in_target_class = False
  self.methods: dict[str, dict] = {}

 def visit_ClassDef(self, node: ast.ClassDef) -> None:
  if node.name == self.class_name:
   self.in_target_class = True
   self.generic_visit(node)
   self.in_target_class = False

 def visit_FunctionDef(self, node: ast.FunctionDef) -> None:
  if not self.in_target_class:
   return

  if not node.name.startswith(self.method_prefix):
   return

  if node.name == "__init__" or node.name.startswith("__"):
   return

  args = [arg.arg for arg in node.args.args]
  self.methods[node.name] = {
   "line": node.lineno,
   "end_line": node.end_lineno,
   "args": args,
   "has_return": any(isinstance(n, ast.Return) for n in ast.walk(node)),
  }


def extract_render_methods(source_path: str, class_name: str = "App", method_prefix: str = "_render") -> dict:
 with open(source_path, "r", encoding="utf-8") as f:
  source = f.read()

 tree = ast.parse(source)
 extractor = AppRenderExtractor(class_name, method_prefix)
 extractor.visit(tree)

 return extractor.methods


def transform_file(source_path: str, output_path: Optional[str] = None, class_name: str = "App", method_prefix: str = "_render") -> list[str]:
 with open(source_path, "r", encoding="utf-8") as f:
  source = f.read()

 tree = ast.parse(source)

 transformer = RenderMethodTransformer(class_name, method_prefix)
 new_tree = transformer.visit(tree)

 result = ast.unparse(new_tree)

 if output_path:
  with open(output_path, "w", encoding="utf-8", newline="") as f:
   f.write(result)
 else:
  print(result)

 return transformer.transformed_methods


def main():
 parser = argparse.ArgumentParser(description="Transform App._render_xxx methods to module-level functions")
 parser.add_argument("source", help="Source Python file to transform")
 parser.add_argument("-o", "--output", help="Output file (default: stdout)")
 parser.add_argument("-c", "--cls", default="App", help="Class name to transform (default: App)")
 parser.add_argument("-p", "--prefix", default="_render", help="Method prefix to extract (default: _render)")
 parser.add_argument("--extract-only", action="store_true", help="Only extract method info, don't transform")

 args = parser.parse_args()

 if args.extract_only:
  methods = extract_render_methods(args.source, args.cls, args.prefix)
  print(f"Found {len(methods)} methods:")
  for name, info in sorted(methods.items(), key=lambda x: x[1]["line"]):
   print(f"  {name}: lines {info['line']}-{info['end_line']}, args={info['args']}")
 else:
  transformed = transform_file(args.source, args.output, args.cls, args.prefix)
  print(f"Transformed {len(transformed)} methods: {transformed}")


if __name__ == "__main__":
 main()