#!/usr/bin/env python3
# Copyright (C) 2025 Riften Labs AS,
# This software is licensed under the GNU Affero General Public License (AGPL), version 3.0 or later.,
# A copy of the license can be found in the LICENSE file or at https://www.gnu.org/licenses/agpl-3.0.html

"""
API Documentation Generator for Riften Labs Indexer

This tool parses routes from main.rs and generates API documentation
in markdown format for each mount point.
"""

import re
import os
import ast
from pathlib import Path
from typing import List, Optional
from dataclasses import dataclass


@dataclass
class RouteInfo:
    """Information about a route"""
    path: str
    function_name: str
    method: str = "GET"
    parameters: List[str] = None
    query_params: List[str] = None
    path_params: List[str] = None
    description: str = ""
    status: str = "unstable"
    is_deprecated: bool = False


@dataclass
class MountInfo:
    """Information about a mount point"""
    path: str
    routes: List[RouteInfo]


class RustRouteParser:
    """Parser for Rust route definitions"""
    
    def __init__(self, main_rs_path: str):
        self.main_rs_path = main_rs_path
        self.mounts: List[MountInfo] = []
        
    def parse_mounts(self) -> List[MountInfo]:
        """Parse all mount points from main.rs"""
        with open(self.main_rs_path, 'r') as f:
            content = f.read()
        
        # Find all mount blocks - simpler approach
        # Look for .mount( followed by path and routes![
        mount_pattern = r'\.mount\(\s*"([^"]+)"\s*,\s*routes!\[(.*?)\]'
        matches = list(re.finditer(mount_pattern, content, re.DOTALL))
        
        for match in matches:
            mount_path = match.group(1)
            routes_content = match.group(2)
            
            routes = self._parse_routes(routes_content)
            mount_info = MountInfo(path=mount_path, routes=routes)
            self.mounts.append(mount_info)
        
        return self.mounts
    
    def _parse_routes(self, routes_content: str) -> List[RouteInfo]:
        """Parse individual routes from a routes! macro"""
        routes = []
        
        # Split by commas and clean up
        route_items = [item.strip() for item in routes_content.split(',') if item.strip()]
        
        for item in route_items:
            # Remove any trailing commas and whitespace
            item = item.strip().rstrip(',')
            if item:
                route_info = self._parse_single_route(item)
                if route_info:
                    routes.append(route_info)
        
        return routes
    
    def _parse_single_route(self, route_item: str) -> Optional[RouteInfo]:
        """Parse a single route item"""
        # Extract function name (remove module prefix if present)
        if '::' in route_item:
            function_name = route_item.split('::')[-1]
        else:
            function_name = route_item
        
        # For now, we'll get the actual route path and parameters from the function definition
        # This is a simplified approach - in a real implementation you'd want to parse the actual function signatures
        return RouteInfo(
            path="",  # Will be filled in by function signature parsing
            function_name=function_name,
            parameters=[],
            query_params=[],
            path_params=[]
        )


class FunctionSignatureParser:
    """Parser for Rust function signatures to extract route information"""
    
    def __init__(self, src_dir: str):
        self.src_dir = src_dir
        self.rpc_modules = {}
        self._load_rpc_modules()
    
    def _load_rpc_modules(self):
        """Load all RPC module files"""
        rpc_dir = os.path.join(self.src_dir, 'rpc')
        if not os.path.exists(rpc_dir):
            return
        
        # Load files directly in rpc directory
        for file in os.listdir(rpc_dir):
            if file.endswith('.rs') and file != 'mod.rs':
                module_name = file[:-3]  # Remove .rs extension
                file_path = os.path.join(rpc_dir, file)
                self.rpc_modules[module_name] = file_path
        
        # Load files in subdirectories
        for subdir in os.listdir(rpc_dir):
            subdir_path = os.path.join(rpc_dir, subdir)
            if os.path.isdir(subdir_path):
                # Look for mod.rs files in subdirectories
                mod_rs_path = os.path.join(subdir_path, 'mod.rs')
                if os.path.exists(mod_rs_path):
                    self.rpc_modules[subdir] = mod_rs_path
    
    def get_function_signature(self, function_name: str) -> Optional[RouteInfo]:
        """Get function signature and route information for a given function name"""
        for module_name, file_path in self.rpc_modules.items():
            route_info = self._parse_function_in_file(file_path, function_name)
            if route_info:
                return route_info
        return None
    
    def _parse_function_in_file(self, file_path: str, function_name: str) -> Optional[RouteInfo]:
        """Parse a specific function in a file"""
        try:
            with open(file_path, 'r') as f:
                content = f.read()
            
            # Look for the function definition
            func_pattern = rf'pub fn {function_name}\s*\([^)]*\)\s*->[^{{]*{{'
            func_match = re.search(func_pattern, content)
            
            if not func_match:
                return None
            
            # Extract description, status, and deprecated flag from comments above the function
            description, status, is_deprecated = self._extract_function_description(content, function_name)
            
            # Skip deprecated functions
            if is_deprecated:
                return None
            
            # Find the #[get("...")] attribute before the function
            get_attr_pattern = rf'#\[get\("([^"]+)"\)\]\s*pub fn {function_name}'
            get_match = re.search(get_attr_pattern, content)
            
            if get_match:
                route_path = get_match.group(1)
                route_info = self._parse_route_path(route_path, function_name)
                route_info.description = description
                route_info.status = status
                route_info.is_deprecated = is_deprecated
                return route_info
            
            # If no #[get] attribute found, return None - we can't determine the route
            return None
            
        except Exception as e:
            print(f"Error parsing {file_path}: {e}")
            return None
    
    def _extract_function_description(self, content: str, function_name: str) -> tuple[str, str, bool]:
        """Extract description, status, and deprecated flag from comments above the function"""
        # Look for the function definition - simpler pattern
        func_pattern = rf'pub fn {function_name}'
        func_match = re.search(func_pattern, content)
        
        if not func_match:
            return "no description", "unstable", False
        
        # Get the position of the function
        func_start = func_match.start()
        
        # Look backwards from the function for comments
        # Find consecutive /// comment lines before the function
        lines = content[:func_start].split('\n')
        comment_lines = []
        
        # Go backwards from the function to find comment lines
        for line in reversed(lines):
            line = line.strip()
            if line.startswith('///'):
                comment_lines.insert(0, line[3:].strip())  # Remove /// and strip
            elif line.startswith('pub fn'):
                break
            elif line and not line.startswith('///') and not line.startswith('#['):
                break
        
        if comment_lines:
            # Filter out status lines and clean up the description
            filtered_lines = []
            for line in comment_lines:
                # Skip status lines (they're handled separately)
                if "Status:" in line:
                    continue
                # Keep empty lines to preserve paragraph structure
                if not line.strip():
                    filtered_lines.append("")
                else:
                    # Preserve indentation by removing only the /// prefix
                    cleaned_line = line[3:] if line.startswith('///') else line
                    filtered_lines.append(cleaned_line.rstrip())
            
            # Join filtered lines, preserving original structure
            description = '\n'.join(filtered_lines).strip()
            
            # Check for deprecated functions
            is_deprecated = any('deprecated' in line.lower() for line in comment_lines)
            
            # Check for status
            status = "unstable"  # default
            for line in comment_lines:
                if "Status: Stable" in line:
                    status = "stable"
                    break
            
            return description, status, is_deprecated
        
        return "no description", "unstable", False

    
    def _parse_route_path(self, route_path: str, function_name: str) -> RouteInfo:
        """Parse route path to extract parameters"""
        path_params = []
        query_params = []
        
        # Split path and query parts
        if '?' in route_path:
            path_part, query_part = route_path.split('?', 1)
        else:
            path_part = route_path
            query_part = ""
        
        # Extract path parameters (e.g., <token_id>)
        path_param_pattern = r'<([^>]+)>'
        path_params = re.findall(path_param_pattern, path_part)
        
        # Extract query parameters (e.g., <token_id>&<timestamp>)
        if query_part:
            query_param_pattern = r'<([^>]+)>'
            query_params = re.findall(query_param_pattern, query_part)
        
        # Clean up the route path for display
        display_path = path_part
        
        return RouteInfo(
            path=display_path,
            function_name=function_name,
            method="GET",
            path_params=path_params,
            query_params=query_params,
            parameters=path_params + query_params
        )


class APIDocGenerator:
    """Generate API documentation in markdown format"""
    
    def __init__(self, output_dir: str = "apidoc"):
        self.output_dir = output_dir
        os.makedirs(output_dir, exist_ok=True)
    
    def generate_docs(self, mounts: List[MountInfo], function_parser: FunctionSignatureParser):
        """Generate documentation for all mounts"""
        # Generate index page
        self._generate_index(mounts)
        
        # Generate individual mount pages
        for mount in mounts:
            self._generate_mount_page(mount, function_parser)
    
    def _generate_index(self, mounts: List[MountInfo]):
        """Generate the main index page"""
        index_content = """# Riften Labs Indexer API Documentation

This document provides API documentation for the Riften Labs Indexer service.

Note that endpoints marked at unsable are subject to change or may be removed.
If you are using unstable endpoints, let us know in the telegram channel so that we can prioritize them.

## Available Endpoints

"""
        
        for mount in mounts:
            mount_name = mount.path.strip('/').replace('/', '_') or 'root'
            mount_display = mount.path if mount.path else '/'
            index_content += f"- [{mount_display}]({mount_name}.md) - {len(mount.routes)} endpoints\n"
        
        index_content += """

## Getting Started

All endpoints return JSON responses. Most endpoints support GET requests only.

### Authentication

Currently, no authentication is required for API access.

"""
        
        with open(os.path.join(self.output_dir, 'index.md'), 'w') as f:
            f.write(index_content)
    
    def _generate_mount_page(self, mount: MountInfo, function_parser: FunctionSignatureParser):
        """Generate documentation for a specific mount point"""
        mount_name = mount.path.strip('/').replace('/', '_') or 'root'
        mount_display = mount.path if mount.path else '/'
        
        content = f"""# {mount_display} Endpoints

Base URL: `{mount_display}`

## Available Endpoints

"""
        
        # Get detailed route information
        detailed_routes = []
        for route in mount.routes:
            detailed_route = function_parser.get_function_signature(route.function_name)
            if detailed_route:
                detailed_routes.append(detailed_route)
            # Skip routes without #[get] attributes - they can't be documented properly
        
        # Sort routes alphabetically by their path
        detailed_routes.sort(key=lambda x: x.path)
        
        for route in detailed_routes:
            content += self._generate_route_doc(route)
        
        # Add navigation
        content += f"""

---

[← Back to Index](index.md)
"""
        
        with open(os.path.join(self.output_dir, f'{mount_name}.md'), 'w') as f:
            f.write(content)
    
    def _generate_route_doc(self, route: RouteInfo) -> str:
        """Generate documentation for a single route"""
        full_path = route.path if route.path.startswith('/') else f'/{route.path}'
        
        # Build the complete route including query parameters
        complete_route = full_path
        if route.query_params:
            complete_route += "?" + "&".join(f"<{param}>" for param in route.query_params)
        
        doc = f"""### `{complete_route}`

**Endpoint:** `{route.method} {full_path}`

**Status:** {route.status.title()}

"""
        
        if route.description and route.description != "no description":
            doc += f"**Description:** {route.description}\n\n"
        
        doc += "---\n\n"
        
        return doc
    

    

    



def main():
    """Main function to generate API documentation"""
    # Configuration
    main_rs_path = "src/main.rs"
    src_dir = "src"
    output_dir = "apidoc"
    
    print("🔍 Parsing routes from main.rs...")
    
    # Parse mounts from main.rs
    route_parser = RustRouteParser(main_rs_path)
    mounts = route_parser.parse_mounts()
    
    print(f"📁 Found {len(mounts)} mount points:")
    for mount in mounts:
        print(f"  - {mount.path}: {len(mount.routes)} routes")
    
    # Initialize function signature parser
    function_parser = FunctionSignatureParser(src_dir)
    
    print("📝 Generating API documentation...")
    
    # Generate documentation
    doc_generator = APIDocGenerator(output_dir)
    doc_generator.generate_docs(mounts, function_parser)
    
    print(f"✅ API documentation generated in '{output_dir}/' directory")
    print("📄 Generated files:")
    
    # List generated files
    for file in os.listdir(output_dir):
        if file.endswith('.md'):
            print(f"  - {output_dir}/{file}")


if __name__ == "__main__":
    main()