from __future__ import annotations

import argparse
import ast
import json
from dataclasses import dataclass
from pathlib import Path
import re
import astor

from route_extractor import (
    RouteInfo,
    extract_routes_from_content
)


@dataclass
class CategoryInfo:
    """Information about a route category"""
    group: str
    patterns: list[str]


def extract_imports(source_code: str) -> list[str]:
    """Extract import statements from source code"""
    tree = ast.parse(source_code)
    imports: set[str] = set()

    for node in ast.walk(tree):
        if isinstance(node, (ast.Import, ast.ImportFrom)):
            imports.add(astor.to_source(node).strip())

    return sorted(imports)

def create_module_content(routes: list[RouteInfo], imports: list[str], category: str) -> tuple[str, list[str]]:
    """
    Create module content with imports and routes.

    Args:
        routes: List of RouteInfo objects
        imports: List of import statements to include
        category: Category name for the routes (e.g. 'user_auth', 'subscription')

    Returns:
        tuple: (file_content, handler_names)
            - file_content: String containing the complete module code
            - handler_names: List of function names to export
    """
    content: list[str] = []
    handler_names: list[str] = []

    # Track function definitions we've already written
    seen_functions: set[str] = set()

    # Add docstring using the category name
    category_display = category.replace('_', ' ').title()
    content.extend([
        f'"""Route handlers for {category_display} endpoints."""',
        ''
    ])

    # Add standard imports
    content.extend([
        'from flask import Response, g, request',
        'from owsresponse import response',
        'from owsresponse.adaptors.flask import flaskify',
    ])

    # Add specific imports
    content.extend(imports)
    content.extend(['', ''])

    # Add route definitions
    for route in routes:
        if route.function_def and route.function not in seen_functions:
            route_code = astor.to_source(route.function_def)
            content.append(route_code)
            content.append('')
            handler_names.append(route.function)
            seen_functions.add(route.function)

    return '\n'.join(content), handler_names


class RouteOrganizer:
    """Organize Flask routes into separate files"""

    def __init__(self, project_root: Path, api_path: str) -> None:
        """
        Initialize the RouteOrganizer

        Args:
            project_root: Path to the project root directory
            api_path: Name of the API directory containing handlers.py (e.g. 'account')
        """
        self.project_root = Path(project_root)
        self.service_dir = self.project_root / 'services'/ api_path
        self.handlers_path = self.service_dir / 'handlers.py'
        self.routes_dir = self.service_dir / 'handlers'

        # Load route categories from JSON config
        config_path = self.service_dir / f'{api_path}.json'
        if not config_path.exists():
            raise FileNotFoundError(f"Could not find config file at {config_path}")

        with config_path.open() as f:
            config = json.load(f)
            self.route_categories = {
                category: CategoryInfo(
                    group=info['group'],
                    patterns=info['patterns']
                )
                for category, info in config['categories'].items()
            }

        if not self.handlers_path.exists():
            raise FileNotFoundError(f"Could not find handlers.py at {self.handlers_path}")

    def categorize_route(self, route_info: RouteInfo) -> str:
        """Categorize route based on its path"""
        if not route_info.route_path:
            return 'misc'

        for category, category_info in self.route_categories.items():
            if any(re.match(pattern, route_info.route_path) for pattern in category_info.patterns):
                return category

        # Debug output for unmatched routes
        print(f"⚠️  No category match found for route: {route_info.route_path}")
        return 'misc'

    def organize(self) -> None:
        """Main organization logic"""
        # Read and parse source file
        source = self.handlers_path.read_text()
        imports = extract_imports(source)
        routes = extract_routes_from_content(source, str(self.handlers_path))

        # Group routes by category
        routes_by_category: dict[str, list[RouteInfo]] = {}
        all_handlers: list[str] = []
        imports_for_init: list[str] = []

        for route in routes:
            category = self.categorize_route(route)
            if category not in routes_by_category:
                routes_by_category[category] = []
            routes_by_category[category].append(route)

        # Create output directory
        self.routes_dir.mkdir(exist_ok=True)

        # Create files for each category and collect handlers
        for category, routes in routes_by_category.items():
            module_name = f"{category}_handlers"
            filename = self.routes_dir / f"{module_name}.py"

            file_content, handler_names = create_module_content(routes, imports, category)
            filename.write_text(file_content)

            if handler_names:
                imports_for_init.append(f"from .{module_name} import (\n    " +
                                     ",\n    ".join(handler_names) +
                                     "\n)")
                all_handlers.extend(handler_names)

        # Create __init__.py content
        init_content = ['"""Initialize and register all route handlers."""\n',
                       "# Import all handlers"]
        init_content.extend(imports_for_init)
        init_content.extend(['', '__all__ = ['])
        init_content.extend([f"    '{handler}'," for handler in sorted(all_handlers)])
        init_content.append(']')

        # Write __init__.py
        (self.routes_dir / '__init__.py').write_text('\n'.join(init_content))

        # Print summary
        print(f"\n✅ Routes organized successfully!")
        print(f"📁 New routes directory: {self.routes_dir}")
        print("\nOrganized routes by category:")

        # Group categories by their group
        groups_dict: dict[str, list[tuple[str, int]]] = {}
        for category, routes in routes_by_category.items():
            group = self.route_categories[category].group
            if group not in groups_dict:
                groups_dict[group] = []
            groups_dict[group].append((category, len(routes)))

        # Print grouped summary
        for group_name, categories in sorted(groups_dict.items()):
            print(f"\n{group_name}:")
            for category, route_count in sorted(categories):
                print(f"  - {category}_handlers.py ({route_count} routes)")


def main() -> None:
    """Main entry point"""
    parser = argparse.ArgumentParser(description="Organize API routes for a service")
    parser.add_argument(
        "--service",
        "-s",
        required=True,
        help="Service name (e.g. 'account', 'permissions')",
    )

    try:
        args = parser.parse_args()
        # Get the project root directory where <service>/ is located
        # Assuming script is run from ows-* directory
        project_root = Path.cwd()
        organizer = RouteOrganizer(project_root=project_root, api_path=args.service)
        organizer.organize()
    except Exception as e:
        print(f"❌ Error: {e}")
        raise


if __name__ == '__main__':
    main()