import ast
import json
import logging
import os
from typing import Any, Dict, List, Optional

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')


class CodeDocumentationExtractor:
    """
    Extracts docstrings, comments, type hints, decorators, and imports from Python files
    and outputs structured JSON.
    """

    def __init__(self) -> None:
        """
        Initializes the CodeDocumentationExtractor.
        """
        pass

    def extract_documentation(self, file_path: str) -> Dict[str, Any]:
        """
        Extracts documentation from a Python file.

        Args:
            file_path (str): The path to the Python file.

        Returns:
            Dict[str, Any]: A dictionary containing the extracted documentation.
        """
        try:
            with open(file_path, 'r', encoding='utf-8') as file:
                source_code = file.read()

            tree = ast.parse(source_code)
            extractor = DocumentationVisitor()
            extractor.visit(tree)

            return {
                'module_docstring': extractor.module_docstring,
                'classes': extractor.classes,
                'functions': extractor.functions,
                'imports': extractor.imports,
            }
        except FileNotFoundError as e:
            logging.error(f"File not found: {file_path} - {e}")
            return {'error': f"File not found: {file_path}"}
        except Exception as e:
            logging.error(f"Error processing file: {file_path} - {e}")
            return {'error': f"Error processing file: {file_path} - {str(e)}"}

    def output_json(self, data: Dict[str, Any], output_path: str) -> None:
        """
        Outputs the extracted documentation to a JSON file.

        Args:
            data (Dict[str, Any]): The extracted documentation.
            output_path (str): The path to the output JSON file.
        """
        try:
            with open(output_path, 'w', encoding='utf-8') as outfile:
                json.dump(data, outfile, indent=4)
            logging.info(f"Documentation saved to: {output_path}")
        except Exception as e:
            logging.error(f"Error writing to file: {output_path} - {e}")


class DocumentationVisitor(ast.NodeVisitor):
    """
    A visitor class that traverses the AST and extracts documentation elements.
    """

    def __init__(self) -> None:
        """
        Initializes the DocumentationVisitor.
        """
        self.module_docstring: Optional[str] = None
        self.classes: List[Dict[str, Any]] = []
        self.functions: List[Dict[str, Any]] = []
        self.imports: List[str] = []

    def visit_Module(self, node: ast.Module) -> None:
        """
        Visits the module node and extracts the module docstring.

        Args:
            node (ast.Module): The module node.
        """
        self.module_docstring = ast.get_docstring(node)
        self.generic_visit(node)

    def visit_ClassDef(self, node: ast.ClassDef) -> None:
        """
        Visits a class definition node and extracts class documentation.

        Args:
            node (ast.ClassDef): The class definition node.
        """
        class_data: Dict[str, Any] = {
            'name': node.name,
            'docstring': ast.get_docstring(node),
            'decorators': [ast.unparse(decorator) for decorator in node.decorator_list],
            'methods': [],
        }

        for body_node in node.body:
            if isinstance(body_node, ast.FunctionDef):
                class_data['methods'].append(self._extract_function_data(body_node))

        self.classes.append(class_data)
        self.generic_visit(node)

    def visit_FunctionDef(self, node: ast.FunctionDef) -> None:
        """
        Visits a function definition node and extracts function documentation.

        Args:
            node (ast.FunctionDef): The function definition node.
        """
        function_data = self._extract_function_data(node)
        self.functions.append(function_data)
        self.generic_visit(node)

    def _extract_function_data(self, node: ast.FunctionDef) -> Dict[str, Any]:
        """
        Extracts function data (docstring, arguments, return type, etc.).

        Args:
            node (ast.FunctionDef): The function definition node.

        Returns:
            Dict[str, Any]: A dictionary containing the extracted function data.
        """
        args: List[Dict[str, Any]] = []
        for arg in node.args.args:
            arg_data: Dict[str, Any] = {
                'name': arg.arg,
                'type_hint': ast.unparse(arg.annotation) if arg.annotation else None
            }
            args.append(arg_data)

        return {
            'name': node.name,
            'docstring': ast.get_docstring(node),
            'args': args,
            'return_type_hint': ast.unparse(node.returns) if node.returns else None,
            'decorators': [ast.unparse(decorator) for decorator in node.decorator_list],
        }

    def visit_Import(self, node: ast.Import) -> None:
        """
        Visits an import statement and extracts the imported modules.

        Args:
            node (ast.Import): The import node.
        """
        for alias in node.names:
            self.imports.append(alias.name)

    def visit_ImportFrom(self, node: ast.ImportFrom) -> None:
        """
        Visits an import from statement and extracts the imported modules.

        Args:
            node (ast.ImportFrom): The import from node.
        """
        for alias in node.names:
            self.imports.append(f"{node.module}.{alias.name}")


if __name__ == '__main__':
    # Example Usage:
    example_file_path = 'example.py'  # Replace with your Python file path
    output_file_path = 'documentation.json'

    # Create a dummy example.py file if it doesn't exist
    if not os.path.exists(example_file_path):
        with open(example_file_path, 'w', encoding='utf-8') as f:
            f.write('''
"""
This is a sample module docstring.
"""

import os
import sys
from typing import List, Optional

@decorator
def example_function(arg1: int, arg2: str) -> bool:
    """
    This is an example function docstring.

    Args:
        arg1: An integer argument.
        arg2: A string argument.

    Returns:
        True if successful, False otherwise.
    """
    return True

class ExampleClass:
    """
    This is an example class docstring.
    """
    def __init__(self, name: str):
        """
        Initializes the ExampleClass.
        """
        self.name = name

    def example_method(self, value: int) -> str:
        """
        This is an example method docstring.

        Args:
            value: An integer value.

        Returns:
            A string representation of the value.
        """
        return str(value)

''')

    extractor = CodeDocumentationExtractor()
    documentation = extractor.extract_documentation(example_file_path)
    extractor.output_json(documentation, output_file_path)

    print(f"Extracted documentation from {example_file_path} and saved to {output_file_path}")