"""
Docstring generation tools for Clover - A terminal assistant for AI-powered project management
"""

import ast
import inspect
import os
import re
import sys
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

# Add the current directory to Python path for imports
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from config.settings import load_config
from models.api_client import APIClient
from tools.file_tools import read_file, update_file


class DocstringGenerator:
    """Handle docstring generation operations with LLM integration"""

    def __init__(self):
        self.config = load_config()
        self.api_client = APIClient()
        self.docstring_styles = {
            "google": self._generate_google_style,
            "numpy": self._generate_numpy_style,
            "sphinx": self._generate_sphinx_style,
            "plain": self._generate_plain_style,
        }

    def _analyze_function_signature(self, node: ast.FunctionDef) -> Dict[str, Any]:
        """
        Analyze function signature to extract parameters and return type

        Args:
            node: AST FunctionDef node

        Returns:
            Dict containing signature analysis
        """
        try:
            # Extract parameters
            params = []
            for arg in node.args.args:
                param_info = {
                    "name": arg.arg,
                    "annotation": None,
                    "default": None,
                }

                # Get type annotation if available
                if arg.annotation:
                    if hasattr(arg.annotation, "id"):
                        param_info["annotation"] = arg.annotation.id
                    else:
                        param_info["annotation"] = ast.unparse(arg.annotation)

                params.append(param_info)

            # Handle defaults
            defaults = node.args.defaults
            if defaults:
                # Defaults apply to the last len(defaults) parameters
                for i, default in enumerate(defaults):
                    param_idx = len(params) - len(defaults) + i
                    if param_idx >= 0 and param_idx < len(params):
                        if hasattr(default, "value"):
                            params[param_idx]["default"] = default.value
                        else:
                            params[param_idx]["default"] = ast.unparse(default)

            # Extract return type annotation
            return_annotation = None
            if node.returns:
                if hasattr(node.returns, "id"):
                    return_annotation = node.returns.id
                else:
                    return_annotation = ast.unparse(node.returns)

            return {
                "name": node.name,
                "parameters": params,
                "return_annotation": return_annotation,
                "is_async": isinstance(node, ast.AsyncFunctionDef),
                "is_method": len(params) > 0 and params[0]["name"] in ["self", "cls"],
                "line_number": node.lineno,
            }

        except Exception as e:
            return {
                "error": f"Error analyzing function signature: {str(e)}",
                "name": node.name,
                "parameters": [],
                "return_annotation": None,
            }

    def _analyze_class_signature(self, node: ast.ClassDef) -> Dict[str, Any]:
        """
        Analyze class signature to extract methods and attributes

        Args:
            node: AST ClassDef node

        Returns:
            Dict containing class analysis
        """
        try:
            methods = []
            attributes = []

            for item in node.body:
                if isinstance(item, ast.FunctionDef):
                    method_info = self._analyze_function_signature(item)
                    methods.append(method_info)
                elif isinstance(item, ast.Assign):
                    # Extract class attributes
                    for target in item.targets:
                        if isinstance(target, ast.Name):
                            attributes.append(target.id)

            # Extract base classes
            bases = []
            for base in node.bases:
                if hasattr(base, "id"):
                    bases.append(base.id)
                else:
                    bases.append(ast.unparse(base))

            return {
                "name": node.name,
                "methods": methods,
                "attributes": attributes,
                "bases": bases,
                "line_number": node.lineno,
            }

        except Exception as e:
            return {
                "error": f"Error analyzing class signature: {str(e)}",
                "name": node.name,
                "methods": [],
                "attributes": [],
            }

    def _generate_google_style(self, signature: Dict[str, Any], purpose: str) -> str:
        """Generate Google-style docstring"""
        lines = [f'"""', purpose, ""]

        if signature.get("parameters"):
            lines.append("Args:")
            for param in signature["parameters"]:
                if param["name"] in ["self", "cls"]:
                    continue

                param_line = f"    {param['name']}"
                if param.get("annotation"):
                    param_line += f" ({param['annotation']})"
                param_line += ": Description of parameter"

                if param.get("default") is not None:
                    param_line += f" (default: {param['default']})"

                lines.append(param_line)
            lines.append("")

        if signature.get("return_annotation"):
            lines.append("Returns:")
            lines.append(
                f"    {signature['return_annotation']}: Description of return value"
            )
        elif not signature.get("is_method") or signature["name"] != "__init__":
            lines.append("Returns:")
            lines.append("    Description of return value")

        lines.append('"""')
        return "\n".join(lines)

    def _generate_numpy_style(self, signature: Dict[str, Any], purpose: str) -> str:
        """Generate NumPy-style docstring"""
        lines = [f'"""', purpose, ""]

        if signature.get("parameters"):
            lines.append("Parameters")
            lines.append("----------")
            for param in signature["parameters"]:
                if param["name"] in ["self", "cls"]:
                    continue

                param_line = param["name"]
                if param.get("annotation"):
                    param_line += f" : {param['annotation']}"
                lines.append(param_line)
                lines.append("    Description of parameter")

                if param.get("default") is not None:
                    lines.append(f"    Default: {param['default']}")
                lines.append("")

        if signature.get("return_annotation"):
            lines.append("Returns")
            lines.append("-------")
            lines.append(f"{signature['return_annotation']}")
            lines.append("    Description of return value")
        elif not signature.get("is_method") or signature["name"] != "__init__":
            lines.append("Returns")
            lines.append("-------")
            lines.append("Description of return value")

        lines.append('"""')
        return "\n".join(lines)

    def _generate_sphinx_style(self, signature: Dict[str, Any], purpose: str) -> str:
        """Generate Sphinx-style docstring"""
        lines = [f'"""', purpose, ""]

        if signature.get("parameters"):
            for param in signature["parameters"]:
                if param["name"] in ["self", "cls"]:
                    continue

                param_line = f":param {param['name']}: Description of parameter"
                if param.get("annotation"):
                    param_line += f"\n:type {param['name']}: {param['annotation']}"
                lines.append(param_line)

        if signature.get("return_annotation"):
            lines.append(f":return: Description of return value")
            lines.append(f":rtype: {signature['return_annotation']}")
        elif not signature.get("is_method") or signature["name"] != "__init__":
            lines.append(":return: Description of return value")

        lines.append('"""')
        return "\n".join(lines)

    def _generate_plain_style(self, signature: Dict[str, Any], purpose: str) -> str:
        """Generate plain docstring"""
        return f'"""{purpose}"""'

    def _generate_docstring_with_llm(
        self, signature: Dict[str, Any], context: str, style: str = "google"
    ) -> str:
        """
        Generate docstring using LLM analysis

        Args:
            signature: Function/class signature information
            context: Surrounding code context
            style: Docstring style to use

        Returns:
            Generated docstring
        """
        try:
            # Prepare prompt for LLM
            if "methods" in signature:  # Class
                prompt = f"""
Generate a comprehensive docstring for the following Python class:

Class name: {signature["name"]}
Base classes: {signature.get("bases", [])}
Methods: {[m["name"] for m in signature.get("methods", [])]}

Context code:
```python
{context}
```

Style: {style}
Requirements:
1. Describe the class purpose and functionality
2. Mention key methods if relevant
3. Follow {style} docstring format
4. Be concise but informative
5. Include usage example if appropriate

Generate only the docstring content (including triple quotes).
"""
            else:  # Function
                params_info = ""
                if signature.get("parameters"):
                    params_info = "Parameters: " + ", ".join(
                        [
                            f"{p['name']}"
                            + (
                                f" ({p.get('annotation', 'Any')})"
                                if p.get("annotation")
                                else ""
                            )
                            for p in signature["parameters"]
                            if p["name"] not in ["self", "cls"]
                        ]
                    )

                return_info = ""
                if signature.get("return_annotation"):
                    return_info = f"Returns: {signature['return_annotation']}"

                prompt = f"""
Generate a comprehensive docstring for the following Python function:

Function name: {signature["name"]}
{params_info}
{return_info}
Is async: {signature.get("is_async", False)}

Context code:
```python
{context}
```

Style: {style}
Requirements:
1. Describe the function purpose and behavior
2. Document all parameters with meaningful descriptions
3. Document return value
4. Follow {style} docstring format
5. Be concise but informative
6. Include usage example if the function is complex

Generate only the docstring content (including triple quotes).
"""

            response = self.api_client.generate_text(
                prompt=prompt, model=self.config.get("model", "qwen2.5-coder:7b")
            )

            if "error" in response:
                # Fallback to template-based generation
                purpose = f"Generated description for {signature['name']}"
                return self.docstring_styles[style](signature, purpose)

            if "choices" in response and len(response["choices"]) > 0:
                content = response["choices"][0]["message"]["content"].strip()
                # Clean up the response - ensure it starts and ends with triple quotes
                if not content.startswith('"""'):
                    content = '"""' + content
                if not content.endswith('"""'):
                    content = content + '"""'
                return content
            else:
                # Fallback
                purpose = f"Generated description for {signature['name']}"
                return self.docstring_styles[style](signature, purpose)

        except Exception as e:
            # Fallback to template generation
            purpose = f"Description for {signature['name']}"
            return self.docstring_styles[style](signature, purpose)


def generate_docstring(
    filepath: str, target: str = None, style: str = "google"
) -> Dict[str, Any]:
    """
    Auto-generate docstrings for functions/classes/modules

    Args:
        filepath (str): Path to Python file
        target (str): Specific function/class name (None for all)
        style (str): Docstring style (google, numpy, sphinx, plain)

    Returns:
        Dict containing generation results
    """
    try:
        if not os.path.exists(filepath):
            return {"error": f"File {filepath} does not exist"}

        if not filepath.endswith(".py"):
            return {"error": f"File {filepath} is not a Python file"}

        content = read_file(filepath)
        tree = ast.parse(content)

        generator = DocstringGenerator()
        results = []

        # Process all functions and classes
        for node in ast.walk(tree):
            if isinstance(node, ast.FunctionDef):
                if target is None or node.name == target:
                    # Check if docstring already exists
                    existing_docstring = ast.get_docstring(node)

                    if existing_docstring is None:
                        # Generate docstring
                        signature = generator._analyze_function_signature(node)

                        # Get context (function definition)
                        lines = content.splitlines()
                        start_line = node.lineno - 1

                        # Find the end of function definition
                        end_line = start_line + 10  # Get some context
                        if end_line >= len(lines):
                            end_line = len(lines) - 1

                        context = "\n".join(lines[start_line : end_line + 1])

                        docstring = generator._generate_docstring_with_llm(
                            signature, context, style
                        )

                        results.append(
                            {
                                "type": "function",
                                "name": node.name,
                                "line": node.lineno,
                                "docstring": docstring,
                                "action": "generated",
                            }
                        )

            elif isinstance(node, ast.ClassDef):
                if target is None or node.name == target:
                    # Check if docstring already exists
                    existing_docstring = ast.get_docstring(node)

                    if existing_docstring is None:
                        # Generate docstring
                        signature = generator._analyze_class_signature(node)

                        # Get context (class definition)
                        lines = content.splitlines()
                        start_line = node.lineno - 1

                        # Find reasonable context for class
                        end_line = start_line + 15  # Get more context for classes
                        if end_line >= len(lines):
                            end_line = len(lines) - 1

                        context = "\n".join(lines[start_line : end_line + 1])

                        docstring = generator._generate_docstring_with_llm(
                            signature, context, style
                        )

                        results.append(
                            {
                                "type": "class",
                                "name": node.name,
                                "line": node.lineno,
                                "docstring": docstring,
                                "action": "generated",
                            }
                        )

        return {
            "filepath": filepath,
            "style": style,
            "target": target,
            "results": results,
            "generated_count": len(results),
        }

    except Exception as e:
        return {"error": f"Error generating docstrings: {str(e)}"}


def update_docstrings(
    filepath: str, target: str = None, style: str = "google"
) -> Dict[str, Any]:
    """
    Update existing docstrings with current function purposes

    Args:
        filepath (str): Path to Python file
        target (str): Specific function/class name (None for all)
        style (str): Docstring style to use

    Returns:
        Dict containing update results
    """
    try:
        if not os.path.exists(filepath):
            return {"error": f"File {filepath} does not exist"}

        content = read_file(filepath)
        tree = ast.parse(content)
        lines = content.splitlines()

        generator = DocstringGenerator()
        results = []
        modifications = []

        # Process all functions and classes with existing docstrings
        for node in ast.walk(tree):
            if isinstance(node, (ast.FunctionDef, ast.ClassDef)):
                if target is None or node.name == target:
                    existing_docstring = ast.get_docstring(node)

                    if existing_docstring is not None:
                        # Generate updated docstring
                        if isinstance(node, ast.FunctionDef):
                            signature = generator._analyze_function_signature(node)
                            node_type = "function"
                        else:
                            signature = generator._analyze_class_signature(node)
                            node_type = "class"

                        # Get context
                        start_line = node.lineno - 1
                        end_line = min(start_line + 15, len(lines) - 1)
                        context = "\n".join(lines[start_line : end_line + 1])

                        new_docstring = generator._generate_docstring_with_llm(
                            signature, context, style
                        )

                        # Find docstring location in source
                        docstring_line = node.lineno  # Line after function/class def

                        # Find the actual docstring lines
                        docstring_start = None
                        docstring_end = None

                        for i in range(
                            docstring_line, min(docstring_line + 10, len(lines))
                        ):
                            line = lines[i].strip()
                            if line.startswith('"""') or line.startswith("'''"):
                                docstring_start = i
                                if line.count('"""') == 2 or line.count("'''") == 2:
                                    # Single line docstring
                                    docstring_end = i
                                else:
                                    # Multi-line docstring - find end
                                    quote = '"""' if line.startswith('"""') else "'''"
                                    for j in range(i + 1, min(i + 20, len(lines))):
                                        if quote in lines[j]:
                                            docstring_end = j
                                            break
                                break

                        if docstring_start is not None and docstring_end is not None:
                            modifications.append(
                                {
                                    "start_line": docstring_start
                                    + 1,  # 1-based for update_file
                                    "end_line": docstring_end + 1,
                                    "new_content": new_docstring,
                                }
                            )

                            results.append(
                                {
                                    "type": node_type,
                                    "name": node.name,
                                    "line": node.lineno,
                                    "old_docstring": existing_docstring,
                                    "new_docstring": new_docstring,
                                    "action": "updated",
                                }
                            )

        # Apply modifications to file
        success_count = 0
        for mod in modifications:
            success = update_file(
                filepath, mod["new_content"], mod["start_line"], mod["end_line"]
            )
            if success:
                success_count += 1

        return {
            "filepath": filepath,
            "style": style,
            "target": target,
            "results": results,
            "updated_count": success_count,
            "total_modifications": len(modifications),
        }

    except Exception as e:
        return {"error": f"Error updating docstrings: {str(e)}"}


def analyze_docstring_coverage(project_path: str = ".") -> Dict[str, Any]:
    """
    Analyze docstring coverage across a project

    Args:
        project_path (str): Path to project directory

    Returns:
        Dict containing coverage analysis
    """
    try:
        python_files = []

        # Find all Python files
        for root, dirs, files in os.walk(project_path):
            # Skip common directories
            dirs[:] = [
                d for d in dirs if d not in {".git", "__pycache__", "venv", "env"}
            ]

            for file in files:
                if file.endswith(".py") and not file.startswith("__"):
                    python_files.append(os.path.join(root, file))

        total_functions = 0
        total_classes = 0
        documented_functions = 0
        documented_classes = 0
        analysis_results = []

        for filepath in python_files:
            try:
                content = read_file(filepath)
                tree = ast.parse(content)

                file_functions = 0
                file_classes = 0
                file_doc_functions = 0
                file_doc_classes = 0

                for node in ast.walk(tree):
                    if isinstance(node, ast.FunctionDef):
                        file_functions += 1
                        total_functions += 1

                        if ast.get_docstring(node):
                            file_doc_functions += 1
                            documented_functions += 1

                    elif isinstance(node, ast.ClassDef):
                        file_classes += 1
                        total_classes += 1

                        if ast.get_docstring(node):
                            file_doc_classes += 1
                            documented_classes += 1

                file_coverage = 0
                if file_functions + file_classes > 0:
                    file_coverage = (
                        (file_doc_functions + file_doc_classes)
                        / (file_functions + file_classes)
                        * 100
                    )

                analysis_results.append(
                    {
                        "filepath": filepath,
                        "functions": file_functions,
                        "classes": file_classes,
                        "documented_functions": file_doc_functions,
                        "documented_classes": file_doc_classes,
                        "coverage_percentage": round(file_coverage, 2),
                    }
                )

            except Exception as e:
                analysis_results.append(
                    {
                        "filepath": filepath,
                        "error": f"Error analyzing file: {str(e)}",
                        "coverage_percentage": 0,
                    }
                )

        # Calculate overall coverage
        total_items = total_functions + total_classes
        documented_items = documented_functions + documented_classes
        overall_coverage = (
            (documented_items / total_items * 100) if total_items > 0 else 0
        )

        return {
            "project_path": project_path,
            "total_files": len(python_files),
            "total_functions": total_functions,
            "total_classes": total_classes,
            "documented_functions": documented_functions,
            "documented_classes": documented_classes,
            "overall_coverage": round(overall_coverage, 2),
            "file_analysis": analysis_results,
        }

    except Exception as e:
        return {"error": f"Error analyzing docstring coverage: {str(e)}"}


def batch_generate_docstrings(
    project_path: str = ".", style: str = "google", overwrite: bool = False
) -> Dict[str, Any]:
    """
    Generate docstrings for all files in a project

    Args:
        project_path (str): Path to project directory
        style (str): Docstring style to use
        overwrite (bool): Whether to overwrite existing docstrings

    Returns:
        Dict containing batch generation results
    """
    try:
        python_files = []

        # Find all Python files
        for root, dirs, files in os.walk(project_path):
            dirs[:] = [
                d for d in dirs if d not in {".git", "__pycache__", "venv", "env"}
            ]

            for file in files:
                if file.endswith(".py") and not file.startswith("__"):
                    python_files.append(os.path.join(root, file))

        results = []
        total_generated = 0

        for filepath in python_files:
            print(f"Processing: {filepath}")

            if overwrite:
                result = update_docstrings(filepath, style=style)
                action = "updated"
            else:
                result = generate_docstring(filepath, style=style)
                action = "generated"

            if "error" not in result:
                count = result.get("generated_count", 0) or result.get(
                    "updated_count", 0
                )
                total_generated += count
                print(f"✓ {action.title()} {count} docstrings in {filepath}")
            else:
                print(f"✗ Error processing {filepath}: {result['error']}")

            results.append(result)

        return {
            "project_path": project_path,
            "total_files": len(python_files),
            "total_generated": total_generated,
            "style": style,
            "overwrite": overwrite,
            "results": results,
        }

    except Exception as e:
        return {"error": f"Error in batch docstring generation: {str(e)}"}


# Example usage
if __name__ == "__main__":
    # Generate docstrings for a specific file
    result = generate_docstring("example.py", style="google")
    print(f"Generated docstrings: {result}")

    # Analyze project coverage
    coverage = analyze_docstring_coverage(".")
    print(f"Docstring coverage: {coverage['overall_coverage']}%")