import ast
import os
import re
from dataclasses import dataclass, field
from typing import List, Dict, Optional

@dataclass
class VariableDoc:
    name: str
    value: str
    docstring: str

@dataclass
class FunctionDoc:
    name: str
    signature: str
    docstring: str

@dataclass
class ModuleDoc:
    name: str
    path: str
    docstring: Optional[str] = None
    functions: List[FunctionDoc] = field(default_factory=list)
    variables: List[VariableDoc] = field(default_factory=list)

class WikinParser:
    def __init__(self, root_dir: str):
        self.root_dir = os.path.abspath(root_dir)
        self.modules: List[ModuleDoc] = []

    def parse(self):
        for root, _, files in os.walk(self.root_dir):
            for file in files:
                if file.endswith(".py"):
                    full_path = os.path.join(root, file)
                    rel_path = os.path.relpath(full_path, self.root_dir)
                    module_name = rel_path.replace(os.sep, ".").replace(".py", "")
                    if module_name.endswith(".__init__"):
                        module_name = module_name[:-9]
                    
                    module_doc = self._parse_file(full_path, module_name)
                    if module_doc.functions or module_doc.variables or module_doc.docstring:
                        self.modules.append(module_doc)
        
        return self.modules

    def _parse_file(self, file_path: str, module_name: str) -> ModuleDoc:
        with open(file_path, "r", encoding="utf-8") as f:
            source = f.read()

        try:
            tree = ast.parse(source)
        except SyntaxError:
            return ModuleDoc(name=module_name, path=file_path)

        module_doc = ModuleDoc(
            name=module_name,
            path=file_path,
            docstring=ast.get_docstring(tree)
        )

        # Extract functions
        for node in tree.body:
            if isinstance(node, ast.FunctionDef):
                docstring = ast.get_docstring(node)
                if docstring:
                    # Get signature
                    args = []
                    for arg in node.args.args:
                        args.append(arg.arg)
                    signature = f"{node.name}({', '.join(args)})"
                    module_doc.functions.append(FunctionDoc(
                        name=node.name,
                        signature=signature,
                        docstring=docstring
                    ))

        # Extract variables with #: docstrings
        # We need to scan the lines for #: comments
        lines = source.splitlines()
        
        # Two types of variable docs:
        # Type 1: #: Doc before\nvar = val
        # Type 2: var = val #: Doc after
        
        for i, line in enumerate(lines):
            stripped = line.strip()
            
            # Skip lines that are just string literals containing #: (like in our own code)
            if 'if "#:" in line' in line or 'line.split("#:", 1)' in line:
                continue

            # Type 2: Post-comment (e.g., var = val #: comment)
            if "#:" in line and "=" in line:
                # Basic check to avoid matching #: inside strings
                # This is a bit naive but covers many cases
                parts = line.split("#:", 1)
                code_part = parts[0].strip()
                comment_part = parts[1].strip()
                
                if "=" in code_part and not code_part.startswith("#"):
                    var_name_part = code_part.split("=", 1)[0].strip()
                    # Check if it's a simple name
                    if re.match(r'^[a-zA-Z_][a-zA-Z0-9_]*$', var_name_part):
                        module_doc.variables.append(VariableDoc(
                            name=var_name_part,
                            value=code_part.split("=", 1)[1].strip(),
                            docstring=comment_part
                        ))
                        continue

            # Type 1: Pre-comment (e.g., #: comment\nvar = val)
            if stripped.startswith("#:"):
                comment = stripped[2:].strip()
                # Look at next line(s) for assignment
                next_idx = i + 1
                while next_idx < len(lines):
                    next_line = lines[next_idx].strip()
                    if not next_line or next_line.startswith("#"):
                        next_idx += 1
                        continue
                    if "=" in next_line:
                        var_name_part = next_line.split("=", 1)[0].strip()
                        if re.match(r'^[a-zA-Z_][a-zA-Z0-9_]*$', var_name_part):
                            module_doc.variables.append(VariableDoc(
                                name=var_name_part,
                                value=next_line.split("=", 1)[1].strip(),
                                docstring=comment
                            ))
                        break
                    else:
                        break

        return module_doc
