import ast
import re
import tempfile
import os
from typing import Dict, List
import pydocstyle


# =====================================================
# PARSER
# =====================================================

def parse_file(file_path):
    with open(file_path, "r") as f:
        source_code = f.read()
    return ast.parse(source_code)


# =====================================================
# EXTRACTION
# =====================================================

def extract_definitions(tree: ast.AST):
    functions = []
    classes = []

    for node in ast.walk(tree):
        if isinstance(node, ast.FunctionDef):
            functions.append(node)
        elif isinstance(node, ast.ClassDef):
            classes.append(node)

    return functions, classes


def get_function_metadata(func: ast.FunctionDef) -> Dict:
    params = [arg.arg for arg in func.args.args]
    docstring = ast.get_docstring(func)
    has_return = any(isinstance(n, ast.Return) for n in ast.walk(func))

    raises = set()
    for n in ast.walk(func):
        if isinstance(n, ast.Raise) and n.exc:
            try:
                exc = ast.unparse(n.exc)
                exc = re.sub(r"\(.*\)", "", exc)
                raises.add(exc)
            except Exception:
                pass

    is_generator = any(
        isinstance(n, (ast.Yield, ast.YieldFrom))
        for n in ast.walk(func)
    )

    return {
        "name": func.name,
        "params": params,
        "has_docstring": docstring is not None,
        "docstring": docstring,
        "has_return": has_return,
        "raises": list(raises),
        "is_generator": is_generator,
    }


# =====================================================
# COVERAGE
# =====================================================

def generate_coverage_report(metadata):
    total_functions = len(metadata)

    documented = [m for m in metadata if m.get("has_docstring")]
    undocumented = [m for m in metadata if not m.get("has_docstring")]

    coverage_percent = (
        round((len(documented) / total_functions) * 100, 2)
        if total_functions > 0 else 0
    )

    return {
        "total_functions": total_functions,
        "documented_functions": len(documented),
        "undocumented_functions": len(undocumented),
        "coverage_percent": coverage_percent
    }


# =====================================================
# DOCSTRING GENERATION
# =====================================================

def generate_docstring(meta: dict, style: str = "google") -> str:
    """
    Generate clean PEP-257 compliant docstring.
    """

    name = meta["name"].replace("_", " ").capitalize()
    params = meta["params"]
    raises = meta["raises"]
    is_generator = meta["is_generator"]
    has_return = meta["has_return"]

    lines = [f"{name}.", ""]

    # ---------------- PARAMETERS ----------------
    if params:
        lines.append("Args:")
        for p in params:
            lines.append(f"    {p}: Description.")
        lines.append("")

    # ---------------- RETURNS / YIELDS ----------------
    if is_generator:
        lines.append("Yields:")
        lines.append("    value: Description.")
        lines.append("")
    elif has_return:
        lines.append("Returns:")
        lines.append("    value: Description.")
        lines.append("")

    # ---------------- RAISES ----------------
    if raises:
        lines.append("Raises:")
        for r in raises:
            lines.append(f"    {r}: Description.")
        lines.append("")

    # Remove trailing blank lines
    while lines and lines[-1] == "":
        lines.pop()

    return "\n".join(lines)


# =====================================================
# INSTRUMENTATION
# =====================================================

def instrument_ast(tree: ast.AST, style: str = "google") -> ast.AST:
    """
    Insert generated docstrings into module, classes, and functions.
    """

    # ---------------- MODULE DOCSTRING ----------------
    if not ast.get_docstring(tree):
        module_doc = ast.Expr(
            value=ast.Constant("Module description.")
        )
        tree.body.insert(0, module_doc)

    for node in ast.walk(tree):

        # ---------------- CLASS DOCSTRING ----------------
        if isinstance(node, ast.ClassDef):
            if not ast.get_docstring(node):
                class_doc = ast.Expr(
                    value=ast.Constant(f"{node.name} class.")
                )
                node.body.insert(0, class_doc)

        # ---------------- FUNCTION DOCSTRING ----------------
        if isinstance(node, ast.FunctionDef):
            if not ast.get_docstring(node):
                meta = get_function_metadata(node)
                docstring = generate_docstring(meta, style)

                node.body.insert(
                    0,
                    ast.Expr(value=ast.Constant(docstring))
                )

    return tree


def instrument_tree(tree, metadata=None, style="google"):
    return instrument_ast(tree, style)


def generate_instrumented_code(tree):
    return ast.unparse(tree)


# =====================================================
# VALIDATOR
# =====================================================

IGNORED_RULES = [
    "D100", "D101", "D102",
    "D205", "D400", "D401", "D212", "D213",
    "D406", "D407", "D413",
]


def validate_file(source_code: str) -> List[Dict]:
    with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as tmp:
        tmp.write(source_code)
        tmp_path = tmp.name

    violations = []

    try:
        for error in pydocstyle.check([tmp_path], ignore=IGNORED_RULES):
            violations.append({
                "line": error.line,
                "code": error.code,
                "message": error.message,
                "source": error.source,
            })
    finally:
        os.unlink(tmp_path)

    return violations
