"""
Code Documentation Generator Utility for Code Assistant Bot
Author: RSK World (https://rskworld.in)
Founder: Molla Samser
Designer & Tester: Rima Khatun
Contact: help@rskworld.in, +91 93305 39277
Year: 2026
"""

import re
import ast
from typing import Dict, List, Any, Optional

class CodeDocumenter:
    """
    Utility class for generating code documentation
    """
    
    def __init__(self):
        self.supported_languages = [
            'python', 'javascript', 'typescript', 'java', 'cpp', 'c',
            'csharp', 'php', 'ruby', 'go', 'rust'
        ]
    
    def generate_documentation(self, code: str, language: str, doc_style: str = 'standard') -> Dict[str, Any]:
        """
        Generate documentation for code
        """
        try:
            if language not in self.supported_languages:
                return {
                    'success': False,
                    'error': f'Language {language} not supported for documentation',
                    'documented_code': code
                }
            
            if language == 'python':
                return self._document_python(code, doc_style)
            elif language in ['javascript', 'typescript']:
                return self._document_javascript(code, doc_style)
            elif language == 'java':
                return self._document_java(code, doc_style)
            else:
                return self._document_generic(code, language, doc_style)
                
        except Exception as e:
            return {
                'success': False,
                'error': f'Documentation generation failed: {str(e)}',
                'documented_code': code
            }
    
    def _document_python(self, code: str, doc_style: str) -> Dict[str, Any]:
        """Generate Python documentation"""
        try:
            documented_code = code
            changes = []
            
            # Parse the code
            try:
                tree = ast.parse(code)
                
                # Find functions and classes without docstrings
                for node in ast.walk(tree):
                    if isinstance(node, ast.FunctionDef):
                        if not ast.get_docstring(node):
                            # Generate docstring
                            func_name = node.name
                            params = [arg.arg for arg in node.args.args]
                            
                            docstring = f'    """\n'
                            docstring += f'    {self._generate_function_description(func_name, params)}\n'
                            docstring += f'\n'
                            
                            if params:
                                docstring += f'    Args:\n'
                                for param in params:
                                    if param != 'self':
                                        docstring += f'        {param}: Description of {param}\n'
                            
                            # Check for return statement
                            has_return = any(isinstance(n, ast.Return) for n in ast.walk(node))
                            if has_return:
                                docstring += f'\n    Returns:\n'
                                docstring += f'        Description of return value\n'
                            
                            docstring += f'    """'
                            
                            # Find the line after function definition
                            func_line = node.lineno - 1
                            lines = documented_code.split('\n')
                            
                            # Insert docstring after function definition
                            if func_line < len(lines):
                                indent = len(lines[func_line]) - len(lines[func_line].lstrip())
                                docstring_lines = docstring.split('\n')
                                indented_docstring = '\n'.join([' ' * indent + line.lstrip() for line in docstring_lines])
                                
                                lines.insert(func_line + 1, indented_docstring)
                                documented_code = '\n'.join(lines)
                                changes.append(f'Added docstring to function {func_name}')
                    
                    elif isinstance(node, ast.ClassDef):
                        if not ast.get_docstring(node):
                            # Generate class docstring
                            class_name = node.name
                            docstring = f'    """\n'
                            docstring += f'    {self._generate_class_description(class_name)}\n'
                            docstring += f'    """'
                            
                            class_line = node.lineno - 1
                            lines = documented_code.split('\n')
                            
                            if class_line >= 0 and class_line < len(lines):
                                indent = len(lines[class_line]) - len(lines[class_line].lstrip())
                                docstring_lines = docstring.split('\n')
                                indented_docstring = '\n'.join([' ' * indent + line.lstrip() for line in docstring_lines])
                                
                                lines.insert(class_line + 1, indented_docstring)
                                documented_code = '\n'.join(lines)
                                changes.append(f'Added docstring to class {class_name}')
            
            except SyntaxError:
                # If parsing fails, add basic comments
                documented_code = self._add_basic_comments(code)
                changes.append('Added basic comments (code parsing failed)')
            
            return {
                'success': True,
                'documented_code': documented_code,
                'changes_made': changes if changes else ['Code already has documentation']
            }
            
        except Exception as e:
            return {
                'success': False,
                'error': f'Python documentation failed: {str(e)}',
                'documented_code': code
            }
    
    def _document_javascript(self, code: str, doc_style: str) -> Dict[str, Any]:
        """Generate JavaScript/TypeScript documentation"""
        try:
            documented_code = code
            changes = []
            
            # Find functions without JSDoc comments
            function_pattern = r'(function\s+(\w+)\s*\([^)]*\)|const\s+(\w+)\s*=\s*(?:\([^)]*\)|function\s*\([^)]*\))\s*=>)'
            
            def add_jsdoc(match):
                func_match = match.group(0)
                func_name = match.group(2) if match.lastindex >= 2 and match.group(2) else (match.group(3) if match.lastindex >= 3 and match.group(3) else 'function')
                
                # Extract parameters
                params_match = re.search(r'\(([^)]*)\)', func_match)
                params = [p.strip().split('=')[0].strip() for p in params_match.group(1).split(',') if p.strip()] if params_match and params_match.group(1) else []
                
                jsdoc = '/**\n'
                jsdoc += f' * {self._generate_function_description(func_name, params)}\n'
                
                if params:
                    for param in params:
                        if param:
                            jsdoc += f' * @param {{type}} {param} - Description of {param}\n'
                
                jsdoc += ' * @returns {type} Description of return value\n'
                jsdoc += ' */\n'
                
                changes.append(f'Added JSDoc to function {func_name}')
                return jsdoc + func_match
            
            documented_code = re.sub(function_pattern, add_jsdoc, documented_code)
            
            return {
                'success': True,
                'documented_code': documented_code,
                'changes_made': changes if changes else ['Code already has documentation']
            }
            
        except Exception as e:
            return {
                'success': False,
                'error': f'JavaScript documentation failed: {str(e)}',
                'documented_code': code
            }
    
    def _document_java(self, code: str, doc_style: str) -> Dict[str, Any]:
        """Generate Java documentation"""
        try:
            documented_code = code
            changes = []
            
            # Find methods without JavaDoc comments
            method_pattern = r'(public|private|protected)?\s*(static)?\s*(\w+)\s+(\w+)\s*\([^)]*\)'
            
            def add_javadoc(match):
                method_name = match.group(4)
                return_type = match.group(3)
                
                javadoc = '/**\n'
                javadoc += f' * {self._generate_function_description(method_name, [])}\n'
                javadoc += ' *\n'
                
                # Extract parameters
                params_match = re.search(r'\(([^)]*)\)', match.group(0))
                if params_match:
                    params = [p.strip().split()[1] if len(p.strip().split()) > 1 else p.strip() 
                             for p in params_match.group(1).split(',') if p.strip()]
                    for param in params:
                        javadoc += f' * @param {param} Description of {param}\n'
                
                javadoc += f' * @return Description of return value\n'
                javadoc += ' */\n'
                
                changes.append(f'Added JavaDoc to method {method_name}')
                return javadoc + match.group(0)
            
            documented_code = re.sub(method_pattern, add_javadoc, documented_code)
            
            return {
                'success': True,
                'documented_code': documented_code,
                'changes_made': changes if changes else ['Code already has documentation']
            }
            
        except Exception as e:
            return {
                'success': False,
                'error': f'Java documentation failed: {str(e)}',
                'documented_code': code
            }
    
    def _document_generic(self, code: str, language: str, doc_style: str) -> Dict[str, Any]:
        """Generic documentation for other languages"""
        try:
            documented_code = self._add_basic_comments(code)
            
            return {
                'success': True,
                'documented_code': documented_code,
                'changes_made': ['Added basic comments']
            }
        except Exception as e:
            return {
                'success': False,
                'error': f'Documentation failed: {str(e)}',
                'documented_code': code
            }
    
    def _generate_function_description(self, func_name: str, params: List[str]) -> str:
        """Generate a description for a function"""
        # Simple heuristic-based description generation
        if 'get' in func_name.lower():
            return f'Gets {func_name.replace("get", "").lower()}'
        elif 'set' in func_name.lower():
            return f'Sets {func_name.replace("set", "").lower()}'
        elif 'is' in func_name.lower() or 'has' in func_name.lower():
            return f'Checks if {func_name.lower()}'
        elif 'create' in func_name.lower() or 'make' in func_name.lower():
            return f'Creates {func_name.lower()}'
        elif 'delete' in func_name.lower() or 'remove' in func_name.lower():
            return f'Removes {func_name.lower()}'
        else:
            return f'Performs {func_name} operation'
    
    def _generate_class_description(self, class_name: str) -> str:
        """Generate a description for a class"""
        return f'Represents a {class_name.lower()}'
    
    def _add_basic_comments(self, code: str) -> str:
        """Add basic comments to code"""
        lines = code.split('\n')
        commented_lines = []
        
        for i, line in enumerate(lines):
            stripped = line.strip()
            
            # Add comment before function/class definitions
            if stripped.startswith(('def ', 'class ', 'function ', 'class ')):
                if i > 0 and not lines[i-1].strip().startswith('#'):
                    indent = len(line) - len(line.lstrip())
                    commented_lines.append(' ' * indent + f'# {stripped}')
                else:
                    commented_lines.append(line)
            else:
                commented_lines.append(line)
        
        return '\n'.join(commented_lines)