import re
import os
from pathlib import Path

from docutils import nodes
from docutils.parsers.rst import Directive, directives
from sphinx.application import Sphinx
from sphinx.util.typing import ExtensionMetadata


class GoApiRefDirective(Directive):
    """
    Directive for generating Go API reference documentation.

    Usage:
    .. go-api-ref:: path/to/gofile.go
       :section: gpu
    """

    required_arguments = 1  # Requires one argument: the path to the Go file
    optional_arguments = 0
    has_content = False
    option_spec = {
        "section": directives.unchanged,  # Optional section filter
    }

    def run(self):
        # Get the path to the Go file
        go_file_path = self.arguments[0]
        env = self.state.document.settings.env

        # Get the section filter if provided
        section_filter = self.options.get("section", None)

        # Resolve the path relative to the document
        doc_dir = Path(env.doc2path(env.docname)).parent
        source_path = (doc_dir / go_file_path).resolve()

        # Check if the file exists
        if not source_path.exists():
            msg = f"Go source file not found: {source_path}"
            return [nodes.warning("", nodes.paragraph("", msg))]

        # Parse the Go file and generate documentation
        functions = parse_go_file(str(source_path))

        # Create a container for the API documentation
        container = nodes.container()
        container["classes"].append("go-api-reference")

        # Add the API documentation to the container
        content = generate_rst_content(functions, section_filter)
        self.state_machine.insert_input(content, source=str(source_path))

        return [container]


def parse_go_file(file_path):
    """Parse a Go file and extract function documentation."""
    with open(file_path, "r") as f:
        content = f.read()

    # Pattern to match function documentation and definition
    pattern = r"(\/\/[^\n]*(?:\n\/\/[^\n]*)*)\n\s*func\s+([A-Za-z0-9_]+)\s*\((.*?)\)\s*(\(.*?\)|\w+)\s*\{"
    matches = re.findall(pattern, content, re.DOTALL)

    functions = []
    for match in matches:
        doc_comment = match[0]
        func_name = match[1]
        params = match[2].strip()
        return_type = match[3].strip()

        # Process the comment lines
        doc_lines = []
        for line in doc_comment.split("\n"):
            if line.strip().startswith("//"):
                # Remove the comment marker and one space after it (if present)
                comment_text = line.strip()[2:]
                if comment_text.startswith(" "):
                    comment_text = comment_text[1:]
                doc_lines.append(comment_text)

        # Extract sections from the doc comment
        description = []
        input_params = []
        output_params = []
        example = []

        current_section = "description"

        for line in doc_lines:
            if line.startswith("Input parameter"):
                current_section = "input"
                input_params.append(line)
            elif line.startswith("Output:"):
                current_section = "output"
                output_params.append(line)
            elif line.startswith("Example:"):
                current_section = "example"
                example.append(line)
            elif current_section == "description":
                description.append(line)
            elif current_section == "input":
                input_params.append(line)
            elif current_section == "output":
                output_params.append(line)
            elif current_section == "example":
                example.append(line)

        # Combine description lines into a single line
        desc_text = " ".join([line.strip() for line in description if line.strip()])

        # Combine output lines into a single line
        output_text = " ".join([line.strip() for line in output_params if line.strip()])

        # Determine the section based on function name
        parts = func_name.split("_")
        section = parts[1] if len(parts) > 1 else "other"

        functions.append(
            {
                "name": func_name,
                "params": params,
                "return_type": return_type,
                "description": desc_text,
                "input_params": "\n".join(input_params).strip(),
                "output_params": output_text,
                "example": "\n".join(example).strip(),
                "section": section.lower(),  # Store the section for filtering
            }
        )

    return functions


def generate_rst_content(functions, section_filter=None):
    """Generate reStructuredText content from parsed functions."""
    lines = []

    # Filter functions by section if a filter is provided
    if section_filter:
        section_filter = section_filter.lower()
        functions = [f for f in functions if f["section"] == section_filter]

        if not functions:
            lines.append(f"No functions found in section: {section_filter}")
            return lines

    # Group functions by prefix if no section filter is provided
    if not section_filter:
        # Group functions by prefix (e.g., GO_gpu_, GO_cpu_)
        function_groups = {}
        for func in functions:
            section = func["section"]
            if section not in function_groups:
                function_groups[section] = []
            function_groups[section].append(func)

        # Define the order of sections (GPU first, then CPU, then others)
        section_order = []

        # Add GPU section first if it exists
        if "gpu" in function_groups:
            section_order.append("gpu")

        # Add CPU section next if it exists
        if "cpu" in function_groups:
            section_order.append("cpu")

        # Add all other sections in alphabetical order
        for prefix in sorted(function_groups.keys()):
            if prefix not in ["gpu", "cpu"]:
                section_order.append(prefix)

        # Write each group in the specified order
        for section in section_order:
            funcs = function_groups[section]
            lines.append(f"{section.upper()} Functions")
            lines.append("-" * len(f"{section.upper()} Functions"))
            lines.append("")

            for func in funcs:
                add_function_documentation(lines, func)
    else:
        # If a section filter is provided, just document those functions without section headers
        for func in functions:
            add_function_documentation(lines, func)

    return lines


def add_function_documentation(lines, func):
    """Add documentation for a single function to the lines list."""
    lines.append(func['name'])
    lines.append("~" * len(f"``{func['name']}``"))
    lines.append("")

    # Function signature
    return_type = func["return_type"]
    if return_type.startswith("(") and return_type.endswith(")"):
        return_type = return_type[1:-1]

    lines.append(".. code-block:: go")
    lines.append("")
    lines.append(f"   func {func['name']}({func['params']}) {return_type}")
    lines.append("")

    # Description
    if func["description"]:
        lines.append(func["description"])
        lines.append("")

    # Input parameters
    if func["input_params"]:
        for input_line in func["input_params"].split("\n"):
            lines.append(input_line)
        lines.append("")

    # Output parameters
    if func["output_params"]:
        lines.append(func["output_params"])
        lines.append("")

    # Example
    if func["example"]:
        # Process the example to properly format code blocks
        example_lines = func["example"].split("\n")
        in_code_block = False

        for i, line in enumerate(example_lines):
            stripped_line = line.strip()

            # Check if this is the Example: line
            if stripped_line == "Example:":
                lines.append("Example:")
                continue

            # Check if we're entering a code block
            if (
                not in_code_block
                and i > 0
                and (
                    stripped_line.startswith("import")
                    or stripped_line.startswith("if")
                    or stripped_line.startswith("for")
                )
            ):
                in_code_block = True
                lines.append("")
                lines.append(".. code-block:: go")
                lines.append("")

            # Add the line to the formatted example
            if in_code_block:
                # For code blocks, add indentation
                lines.append(f"   {line}")
            elif stripped_line:  # Only add non-empty lines outside code blocks
                lines.append(line)

        lines.append("")


def setup(app):
    """
    Setup function for Sphinx extension.
    This will be called by Sphinx when the extension is loaded.
    """
    # Register the directive
    app.add_directive("go-api-ref", GoApiRefDirective)

    return {
        "version": "0.1.0",
        "parallel_read_safe": True,
        "parallel_write_safe": True,
    }
