#!/usr/bin/env python3
import os
import re
import ast
import shutil

class FortranParser:
    """A robust line-by-line state-machine parser for Fortran free-form source files.
    Correctly resolves line continuations (`&`) and extracts Doxygen/FORD-style comments.
    """
    def __init__(self):
        # Match MODULE name (excluding MODULE PROCEDURE)
        self.module_re = re.compile(r'^\s*module\s+(\w+)\s*$', re.IGNORECASE)
        # Match SUBROUTINE name
        self.subroutine_re = re.compile(r'^\s*(?:recursive\s+|mpich\s+)?subroutine\s+(\w+)\s*(?:\(([^)]*)\))?', re.IGNORECASE)
        # Match FUNCTION name
        self.function_re = re.compile(
            r'^\s*(?:recursive\s+)?(?:real|integer|double\s+precision|logical|character|complex)?(?:\([\w*=]+\))?\s*function\s+(\w+)\s*(?:\(([^)]*)\))?', 
            re.IGNORECASE
        )

    def parse_file(self, filepath):
        with open(filepath, 'r', encoding='utf-8') as f:
            raw_lines = f.readlines()
        
        # Resolve line continuations using &
        lines = []
        current_line = ""
        for line in raw_lines:
            trimmed = line.strip()
            # If the entire line is a comment, keep it as is
            if trimmed.startswith('!'):
                lines.append((line.rstrip('\r\n'), True))
                continue
                
            # Partition the line by '!' to separate code and inline comments
            code_part = ""
            in_string = False
            str_char = None
            for i, char in enumerate(line):
                if char in ("'", '"'):
                    if not in_string:
                        in_string = True
                        str_char = char
                    elif str_char == char:
                        in_string = False
                elif char == '!' and not in_string:
                    code_part = line[:i]
                    break
            else:
                code_part = line
                
            code_trimmed = code_part.strip()
            if code_trimmed.endswith('&'):
                current_line += code_trimmed[:-1].strip() + " "
            else:
                if current_line:
                    lines.append((current_line + code_trimmed, False))
                    current_line = ""
                else:
                    lines.append((line.rstrip('\r\n'), False))

        # Parse resolved lines
        elements = []
        current_comments = []
        module_doc = ""
        in_module = False
        
        for line_str, is_cmt in lines:
            trimmed = line_str.strip()
            if is_cmt or trimmed.startswith('!'):
                if trimmed.startswith('!>') or trimmed.startswith('!!'):
                    cleaned_comment = trimmed[2:].strip()
                    current_comments.append(cleaned_comment)
                continue
                
            if not trimmed:
                if not in_module and not elements and current_comments:
                    module_doc = "\n".join(current_comments)
                    current_comments = []
                continue
                
            m_mod = self.module_re.match(trimmed)
            m_sub = self.subroutine_re.match(trimmed)
            m_func = self.function_re.match(trimmed)
            
            if m_mod:
                in_module = True
                mod_name = m_mod.group(1)
                elements.append({
                    'type': 'module',
                    'name': mod_name,
                    'doc': "\n".join(current_comments) if current_comments else ""
                })
                current_comments = []
            elif m_sub:
                sub_name = m_sub.group(1)
                args_str = m_sub.group(2) or ""
                args = [a.strip() for a in args_str.split(',') if a.strip()]
                elements.append({
                    'type': 'subroutine',
                    'name': sub_name,
                    'args': args,
                    'doc': "\n".join(current_comments) if current_comments else ""
                })
                current_comments = []
            elif m_func:
                func_name = m_func.group(1)
                args_str = m_func.group(2) or ""
                args = [a.strip() for a in args_str.split(',') if a.strip()]
                elements.append({
                    'type': 'function',
                    'name': func_name,
                    'args': args,
                    'doc': "\n".join(current_comments) if current_comments else ""
                })
                current_comments = []
            else:
                # Clear comment buffer if normal statement is met
                current_comments = []
                
        return {
            'doc': module_doc,
            'elements': elements
        }


class PythonASTParser:
    """Uses Python standard library `ast` module to safely parse classes, methods,
    decorators, functions, and Google-style docstrings from Python modules.
    """
    def get_node_name(self, node):
        if hasattr(ast, 'unparse'):
            return ast.unparse(node).strip()
        if isinstance(node, ast.Name):
            return node.id
        elif isinstance(node, ast.Attribute):
            return f"{self.get_node_name(node.value)}.{node.attr}"
        elif isinstance(node, ast.Call):
            return f"{self.get_node_name(node.func)}(...)"
        return str(node)

    def parse_file(self, filepath):
        with open(filepath, 'r', encoding='utf-8') as f:
            content = f.read()
        try:
            tree = ast.parse(content)
        except SyntaxError as e:
            print(f"Syntax error parsing {filepath}: {e}")
            return None
        
        module_doc = ast.get_docstring(tree) or ""
        elements = []
        
        for node in ast.iter_child_nodes(tree):
            if isinstance(node, ast.ClassDef):
                class_doc = ast.get_docstring(node) or ""
                methods = []
                for child in ast.iter_child_nodes(node):
                    if isinstance(child, ast.FunctionDef):
                        method_doc = ast.get_docstring(child) or ""
                        args = [a.arg for a in child.args.args]
                        decorators = [self.get_node_name(d) for d in child.decorator_list]
                        methods.append({
                            'name': child.name,
                            'args': args,
                            'doc': method_doc,
                            'decorators': decorators
                        })
                elements.append({
                    'type': 'class',
                    'name': node.name,
                    'bases': [self.get_node_name(b) for b in node.bases],
                    'doc': class_doc,
                    'methods': methods
                })
            elif isinstance(node, ast.FunctionDef):
                func_doc = ast.get_docstring(node) or ""
                args = [a.arg for a in node.args.args]
                decorators = [self.get_node_name(d) for d in node.decorator_list]
                elements.append({
                    'type': 'function',
                    'name': node.name,
                    'args': args,
                    'doc': func_doc,
                    'decorators': decorators
                })
        return {
            'doc': module_doc,
            'elements': elements
        }


def format_docstring(doc):
    """Formats raw docstrings/comments into clean Markdown."""
    if not doc:
        return "*No description provided.*"
    
    # Process Google Style or Doxygen tags into neat markdown sections
    lines = doc.split('\n')
    formatted_lines = []
    in_args = False
    in_returns = False
    
    for line in lines:
        trimmed = line.strip()
        
        # Google style section headers
        if trimmed == 'Args:':
            formatted_lines.append("\n**Arguments:**\n")
            in_args = True
            in_returns = False
            continue
        elif trimmed == 'Returns:':
            formatted_lines.append("\n**Returns:**\n")
            in_args = False
            in_returns = True
            continue
        elif trimmed in ('Raises:', 'Attributes:'):
            formatted_lines.append(f"\n**{trimmed[:-1]}:**\n")
            in_args = in_returns = False
            continue
            
        # Parse params inside arguments or return
        if in_args and trimmed:
            # Format "arg_name (type): description"
            match = re.match(r'^([\w_]+)\s*(?:\(([^)]+)\))?\s*:\s*(.*)$', trimmed)
            if match:
                name, arg_type, desc = match.groups()
                type_str = f" *({arg_type})*" if arg_type else ""
                formatted_lines.append(f"- ` {name} `{type_str}: {desc}")
                continue
        
        # Parse Doxygen tags
        if trimmed.startswith('@param'):
            match = re.match(r'^@param\s+([\w_]+)\s*(.*)$', trimmed)
            if match:
                name, desc = match.groups()
                formatted_lines.append(f"- ` {name} `: {desc}")
                continue
        elif trimmed.startswith('@return'):
            desc = trimmed.replace('@return', '').strip()
            formatted_lines.append(f"\n**Returns:**\n\n{desc}")
            continue
        elif trimmed.startswith('@note'):
            desc = trimmed.replace('@note', '').strip()
            formatted_lines.append(f"\n> [!NOTE]\n> {desc}")
            continue
        elif trimmed.startswith('@author'):
            desc = trimmed.replace('@author', '').strip()
            formatted_lines.append(f"\n*Author: {desc}*")
            continue
            
        formatted_lines.append(line)
        
    return "\n".join(formatted_lines).strip()


def build_all_docs():
    workspace_dir = os.path.dirname(os.path.abspath(__file__))
    code_dir = os.path.join(workspace_dir, "code")
    docs_dir = os.path.join(workspace_dir, "docs")

    # Recreate the docs directory
    if os.path.exists(docs_dir):
        shutil.rmtree(docs_dir)
    os.makedirs(docs_dir, exist_ok=True)

    fortran_parser = FortranParser()
    py_parser = PythonASTParser()

    # Files to parse
    fortran_files = [
        "Compact.f90",
        "m_arrays.f90",
        "m_calculate.f90",
        "m_openmpi.f90",
        "m_parameters.f90",
        "post.f90",
        "post_dns.f90"
    ]

    python_files = [
        ("code_gen.py", "code/code_gen/code_gen.py"),
        ("post.py", "code/code_gen/post.py"),
        ("pycompact.py", "code/pycompact/pycompact.py")
    ]

    all_modules = []

    # 1. Parse Fortran files
    for fname in fortran_files:
        path = os.path.join(code_dir, fname)
        if not os.path.exists(path):
            continue
        
        print(f"Parsing Fortran file: {fname}")
        res = fortran_parser.parse_file(path)
        if not res:
            continue
            
        all_modules.append({
            'name': fname,
            'lang': 'fortran',
            'data': res
        })

    # 2. Parse Python files
    for key_name, rel_path in python_files:
        path = os.path.join(workspace_dir, rel_path)
        if not os.path.exists(path):
            continue
            
        print(f"Parsing Python file: {key_name}")
        res = py_parser.parse_file(path)
        if not res:
            continue
            
        all_modules.append({
            'name': key_name,
            'lang': 'python',
            'data': res
        })

    # 3. Generate Markdown files
    for mod in all_modules:
        mname = mod['name']
        lang = mod['lang']
        data = mod['data']
        
        md_filename = mname.replace('.f90', '.md').replace('.py', '.md')
        md_path = os.path.join(docs_dir, md_filename)
        
        with open(md_path, 'w', encoding='utf-8') as md_file:
            # Header banner
            badge = "Fortran Source" if lang == 'fortran' else "Python Source"
            badge_color = "blue" if lang == 'fortran' else "green"
            
            md_file.write(f"# {mname}\n\n")
            md_file.write(f"![Language](https://img.shields.io/badge/Language-{badge}-{badge_color}?style=for-the-badge)\n\n")
            
            # Smart Overview Fallback
            overview_doc = data['doc'].strip()
            if not overview_doc:
                for elem in data['elements']:
                    if elem['type'] == 'module' and elem.get('doc'):
                        overview_doc = elem['doc'].strip()
                        break
                else:
                    for elem in data['elements']:
                        if elem.get('doc'):
                            overview_doc = elem['doc'].strip()
                            break
            
            md_file.write("## Overview\n\n")
            md_file.write(format_docstring(overview_doc) + "\n\n")
            
            # Table of Contents
            md_file.write("## API Index\n\n")
            has_elements = False
            for elem in data['elements']:
                has_elements = True
                elem_type = elem['type'].capitalize()
                md_file.write(f"- [{elem_type}: `{elem['name']}`](#{elem['name'].lower()})\n")
            if not has_elements:
                md_file.write("*No direct class or function definitions extracted from this file.*\n")
            md_file.write("\n---\n\n")
            
            # API Reference details
            md_file.write("## API Reference\n\n")
            for elem in data['elements']:
                el_name = elem['name']
                el_type = elem['type']
                el_doc = format_docstring(elem.get('doc', ''))
                
                # Format headers depending on type
                if el_type == 'module':
                    md_file.write(f"### [Module] `{el_name}`\n\n")
                    md_file.write(el_doc + "\n\n")
                elif el_type == 'class':
                    bases_str = f"({', '.join(elem['bases'])})" if elem['bases'] else ""
                    md_file.write(f"### [Class] `{el_name}{bases_str}`\n\n")
                    md_file.write(el_doc + "\n\n")
                    
                    if elem['methods']:
                        md_file.write("#### Methods\n\n")
                        for method in elem['methods']:
                            m_args = ", ".join(method['args'])
                            decs = "".join([f"`@{d}`<br>" for d in method['decorators']])
                            md_file.write(f"##### `{method['name']}({m_args})`\n\n")
                            if decs:
                                md_file.write(f"**Decorators:** {decs}\n\n")
                            md_file.write(format_docstring(method['doc']) + "\n\n")
                elif el_type in ('subroutine', 'function'):
                    args_str = ", ".join(elem['args'])
                    decs = "".join([f"`@{d}`<br>" for d in elem.get('decorators', [])])
                    elem_badge = "Subroutine" if el_type == 'subroutine' else "Function"
                    
                    md_file.write(f"### [{elem_badge}] `{el_name}`\n\n")
                    md_file.write("```fortran\n" if lang == 'fortran' else "```python\n")
                    md_file.write(f"{el_type} {el_name}({args_str})\n")
                    md_file.write("```\n\n")
                    
                    if decs:
                        md_file.write(f"**Decorators:** {decs}\n\n")
                        
                    md_file.write(el_doc + "\n\n")
                
                md_file.write("---\n\n")

    # 4. Generate master README.md index
    readme_path = os.path.join(docs_dir, "README.md")
    with open(readme_path, 'w', encoding='utf-8') as ri:
        ri.write("# HPC DNS Post-Processing Reference Manual\n\n")
        ri.write("> [!NOTE]\n")
        ri.write("> This reference documentation is automatically generated from codebase comments and AST analysis using the literal programming compiler.\n\n")
        
        ri.write("## Project Architecture Overview\n\n")
        ri.write("This suite compiles algebraic DSL budget equations, performs high-order compact finite difference calculus, and exports physical terms via MPI-IO. Below is the primary software module pipeline:\n\n")
        
        # Include a Mermaid architecture diagram
        ri.write("```mermaid\n")
        ri.write("graph TD\n")
        ri.write("    A[code_gen.py / DSL Input] -->|AST Compiler Pipeline| B[post.py / Code Gen]\n")
        ri.write("    B -->|Generates budget kernels| C[m_calculate.f90]\n")
        ri.write("    C -->|LU Tridiagonal Solvers| D[Compact.f90]\n")
        ri.write("    C -->|Array allocations| E[m_arrays.f90]\n")
        ri.write("    F[post.f90 / Main Driver] -->|Orchestrates calculations| C\n")
        ri.write("    F -->|Parallelizes domain| G[m_openmpi.f90]\n")
        ri.write("    F -->|Loads run configuration| H[m_parameters.f90]\n")
        ri.write("    I[pycompact.py / Bindings] -->|Wraps CPU calculations| C\n")
        ri.write("```\n\n")
        
        ri.write("## Reference Files Index\n\n")
        ri.write("| Document / Code Module | Language | Description | \n")
        ri.write("| --- | --- | --- |\n")
        
        for mod in all_modules:
            mname = mod['name']
            lang_label = "Fortran 90" if mod['lang'] == 'fortran' else "Python"
            doc_link = mname.replace('.f90', '.md').replace('.py', '.md')
            
            # Smart summary from first line of doc
            overview_doc = mod['data']['doc'].strip()
            if not overview_doc:
                for elem in mod['data']['elements']:
                    if elem['type'] == 'module' and elem.get('doc'):
                        overview_doc = elem['doc'].strip()
                        break
                else:
                    for elem in mod['data']['elements']:
                        if elem.get('doc'):
                            overview_doc = elem['doc'].strip()
                            break
            
            doc_lines = overview_doc.split('\n')
            brief = doc_lines[0] if doc_lines and doc_lines[0] else "*No overview description provided.*"
            brief = brief.replace('@brief', '').replace('@author', '').strip()
            if brief.startswith('Google') or brief.startswith('Ignis'):
                for line in doc_lines:
                    line_clean = line.replace('@brief', '').replace('@author', '').strip()
                    if line_clean and not line.strip().startswith('@'):
                        brief = line_clean
                        break
                
            ri.write(f"| [{mname}]({doc_link}) | **{lang_label}** | {brief} |\n")
            
        ri.write("\n---\n")
        ri.write("*Generated dynamically by `build_docs.py`.*")

    print(f"\nDocumentation built successfully under {docs_dir}!")


if __name__ == "__main__":
    build_all_docs()