import os
import ast
import argparse
import logging
from typing import List, Dict, Any, Union

# Configure logger to print to output
logging.basicConfig(level=logging.DEBUG, handlers=[logging.StreamHandler()])
logger = logging.getLogger(__name__)

class PythonClassParser:
    def __init__(self, directory: str, output_file: str, ignored_paths: List[str], recursive: bool) -> None:
        logger.debug(f"Initializing PythonClassParser: {directory} : {output_file}")
        self.directory: str = directory
        self.output_file: str = output_file
        self.printed_module_names: set[str] = set()
        self.ignored_paths: List[str] = ignored_paths
        self.recursive: bool = recursive

    @classmethod
    def is_ignored(cls, path: str, ignored_paths: List[str]) -> bool:
        for ignored in ignored_paths:
            if ignored in path:
                return True
        return False

    def find_python_files(self, directory: str) -> List[str]:
        python_files: List[str] = []
        try:
            for root, dirs, files in os.walk(directory):
                if not self.recursive:
                    dirs[:] = []  # remove subdirectories from search
                for file in files:
                    if file != '__init__.py' and file.endswith('.py'):
                        file_path = os.path.join(root, file)
                        if not self.is_ignored(file_path, self.ignored_paths):
                            python_files.append(file_path)
        except OSError as e:
            logger.error(f"Failed to walk directory {directory}: {e}")
        return python_files

    @classmethod
    def parse_file(cls, file: str) -> Union[ast.Module, None]:
        try:
            with open(file, "r") as source:
                return ast.parse(source.read())
        except OSError as e:
            logger.error(f"Failed to open {file}: {e}")
            return None
        except SyntaxError as e:
            logger.error(f"Failed to parse {file}: {e}")
            return None

    @classmethod
    def get_return_type(cls, node: ast.FunctionDef) -> str:
        try:
            if isinstance(node.returns, ast.Name):
                result = node.returns.id
            elif isinstance(node.returns, ast.Subscript):
                result = ast.unparse(node.returns).strip()
            elif isinstance(node.returns, ast.Attribute):
                result = ast.unparse(node.returns).strip()
            else:
                result = 'None'
                
            return cls.shorten_return_type(result)
        except TypeError as e:
            logger.error(f"Failed to get return type for {node.name}: {e}")
            return 'None'


    @classmethod
    def shorten_return_type(cls, return_type: str) -> str:
        # Turn 'Union[None, str]' into 'str'
        if return_type.startswith('Union[None'):
            return_type = return_type[return_type.rfind(',') + 2:-1]
        
        # Turn 'Optional[str]' into 'str'
        if return_type.startswith('Optional['):
            return_type = return_type[return_type.rfind('[') + 1:-1]
        
        # Turn List[str] into str[]
        if return_type.startswith('List['):
            return_type = return_type.replace('List[', '')[:-1] + '[]'
        
        # Turn Dict[str, int] into {int}
        if return_type.startswith('Dict['):
            return_type = return_type.replace('Dict[', '{')[:-1] + '}'
        
        # Turn Tuple[str, int] into (str, int)
        if return_type.startswith('Tuple['):
            return_type = return_type.replace('Tuple[', '(')[:-1] + ')'
        
        # Turn Union[str, int] into str|int
        if return_type.startswith('Union['):
            return_type = return_type.replace('Union[', '').replace(',', '|')[1:-1]
        
        # Turn Iterable[str] into str[]
        if return_type.startswith('Iterable['):
            return_type = return_type.replace('Iterable[', '')[:-1] + '[]'

        # Turn Any[] into []
        if return_type == 'Any[]':
            return_type = '[]'

        return return_type

    @classmethod
    def get_method_signature(cls, node: ast.FunctionDef) -> str:
        try:
            args = [arg.arg for arg in node.args.args if arg.arg not in ['self', 'cls']]

            method = f"{node.name}({', '.join(args)})"
            return_type = cls.get_return_type(node)
            if return_type == 'None':
                return method
            return f"{method} {cls.get_return_type(node)}"
        except Exception as e:
            logger.error(f"Failed to get method signature for {node.name}: {e}")
            return ''

    @classmethod
    def find_classes_and_methods(cls, tree: ast.Module) -> Dict[str, Dict[str, Any]]:
        classes: Dict[str, Dict[str, Any]] = {}
        for node in ast.walk(tree):
            if isinstance(node, ast.ClassDef):
                class_name: str = node.name
                base_classes: List[str] = [base.id for base in node.bases if isinstance(base, ast.Name)]
                method_names: List[str] = [cls.get_method_signature(n) for n in node.body if isinstance(n, ast.FunctionDef) and not n.name.startswith('_')]
                if method_names:  # Ignore classes with no methods (other than __init__)
                    classes[class_name] = {"name": class_name, "base_classes": base_classes, "methods": method_names}
        return classes
    
    @classmethod
    def group_classes_by_module(cls, classes: Dict[str, List[Dict[str, Any]]], module_name: str) -> List[Dict[str, Any]]:
        if module_name not in classes:
            classes[module_name] = []
        return classes[module_name]

    def write_to_file(self, grouped_classes: Dict[str, List[Dict[str, Any]]]) -> None:
        with open(self.output_file, 'a') as f:
            for module_name, classes in grouped_classes.items():
                if not classes:  # Skip empty modules
                    continue

                f.write(f"#{module_name}\n")

                for class_data in classes:
                    try:
                        base_classes_str = '(' + ', '.join(class_data['base_classes']) + ')' if class_data['base_classes'] else ''
                        f.write(f"{class_data['name']}{base_classes_str}\n")
                        for method in class_data['methods']:
                            f.write(f"\t{method}\n")

                        logger.debug(f"Wrote {class_data['name']} and {len(class_data['methods'])} methods to {self.output_file}")
                    except KeyError as e:
                        logger.error(f"Failed to write {class_data}: {e}")

                f.write("\n")

    def run(self) -> None:
        # Clear output file
        if os.path.exists(self.output_file):
            os.remove(self.output_file)

        # Write header text to output file
        with open(self.output_file, 'a') as f:
            f.write("Summarized as:\n#Directory\nClass\n\tmethods\n\nOmitted self/cls params, methods beginning with _\n\n--\n\n")

        grouped_classes: Dict[str, List[Dict[str, Any]]] = {}
        for file in self.find_python_files(self.directory):
            tree = self.parse_file(file)
            if tree is not None:
                logger.debug(f"Parsing {file}")
                module_classes = self.find_classes_and_methods(tree)
                relative_path = os.path.relpath(file, self.directory)
                module_name = os.path.dirname(relative_path).replace(os.path.sep, '.')
                self.group_classes_by_module(grouped_classes, module_name).extend(module_classes.values())
            else:
                logger.warning(f"Failed to parse {file}")
        self.write_to_file(grouped_classes)
        logger.info(f"Finished parsing {self.directory}")

def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Find all classes and their methods in Python files in a directory.")
    parser.add_argument('-d', '--directory', type=str, default='.', help="Directory to search for Python files")
    parser.add_argument('-o', '--output', type=str, default='project_summary.txt', help="Output file")
    parser.add_argument('-i', '--ignore', nargs='*', default=[], help="Directories or files to ignore")
    parser.add_argument('-l', '--logfile', type=str, help="Log file")
    parser.add_argument('-r', '--recursive', action='store_true', help="Enable or disable recursive search in subdirectories", default=True)
    parser.add_argument('-v', '--verbose', action='store_true', help="Print debug messages")
    return parser.parse_args()

if __name__ == "__main__":
    args: argparse.Namespace = parse_args()
    if args.verbose:
        logger.setLevel(logging.DEBUG)

    if args.logfile:
        file_handler = logging.FileHandler(args.logfile)
        logger.addHandler(file_handler)

    parser: PythonClassParser = PythonClassParser(args.directory, args.output, args.ignore, args.recursive)
    parser.run()