diff --git a/src/core/generator.py b/src/core/generator.py
index c4b3904..bdcdb87 100644
--- a/src/core/generator.py
+++ b/src/core/generator.py
@@ -1,70 +1,71 @@
+"""Generator for documentation output."""
+
 import json
-import os
 from pathlib import Path
-from .parser import parse_openapi_spec, load_spec_file
-from .models import Endpoint, APISpec
+from typing import Any, Optional
+
+from src.core.models import OpenAPISpec
+from src.core.parser import _basic_validate, parse_openapi_spec
 
 
-def generate_docs(spec_file: str, output: str = None, format: str = 'html', open_browser: bool = False):
-    """Generate documentation in the specified format.
+def generate_docs(
+    spec_source: str | Path | dict[str, Any],
+    format: str = "html",
+    output_path: Optional[str] = None,
+    template_path: Optional[str] = None,
+) -> str:
+    if isinstance(spec_source, dict):
+        spec_data = spec_source
+    elif isinstance(spec_source, (str, Path)):
+        spec_path = Path(spec_source)
+        if spec_path.exists():
+            content = spec_path.read_text()
+            if spec_path.suffix in [".yaml", ".yml"]:
+                import yaml
 
-    Args:
-        spec_file: Path to the OpenAPI spec file
-        output: Output file path or directory
-        format: Output format (html, markdown, json, all)
-        open_browser: Whether to open the generated file in browser
-    """
-    result = parse_openapi_spec(spec_file)
+                spec_data = yaml.safe_load(content)
+            else:
+                spec_data = json.loads(content)
+        else:
+            raise FileNotFoundError(f"Spec file not found: {spec_source}")
+    else:
+        raise ValueError(f"Invalid spec source type: {type(spec_source)}")
 
-    if not result.get('valid'):
-        raise ValueError(f"Invalid spec: {result.get('errors')}")
+    is_valid, errors = _basic_validate(spec_data)
+    if not is_valid:
+        raise ValueError(f"Invalid spec: {errors}")
 
-    spec = result['spec']
-
-    if format == 'all':
-        for fmt in ['html', 'markdown', 'json']:
-            generate_docs(spec_file, output, fmt, open_browser and fmt == 'html')
-        return
-
-    if not output:
-        base_name = Path(spec_file).stem
-        if format == 'html':
-            output = f"{base_name}.html"
-        elif format == 'markdown':
-            output = f"{base_name}.md"
-        elif format == 'json':
-            output = f"{base_name}_docs.json"
-
-    if format == 'html':
-        from .templates.html_template import generate_html
-        generate_html(spec, output)
-    elif format == 'markdown':
-        from .templates.markdown_template import generate_markdown
-        generate_markdown(spec, output)
-    elif format == 'json':
-        generate_json_docs(spec, output)
-
-    print(f"Generated {format} documentation: {output}")
-
-    if open_browser and format == 'html':
-        import webbrowser
-        webbrowser.open(f'file://{os.path.abspath(output)}')
+    spec = parse_openapi_spec(spec_data)
+    return spec
 
 
-def generate_json_docs(spec: dict, output: str):
-    """Generate JSON documentation."""
-    from .parser import extract_endpoints
+def extract_endpoints(spec: OpenAPISpec) -> list:
+    endpoints = []
+    for path, path_item in spec.paths.items():
+        for method, operation in path_item.model_dump().items():
+            if method in ["get", "put", "post", "delete", "options", "head", "patch", "trace"]:
+                if operation:
+                    endpoints.append(
+                        {
+                            "path": path,
+                            "method": method.upper(),
+                            "summary": operation.get("summary", ""),
+                            "description": operation.get("description", ""),
+                            "tags": operation.get("tags", []),
+                        }
+                    )
+    return endpoints
 
-    endpoints = extract_endpoints(spec)
 
-    docs = {
-        'title': spec.get('info', {}).get('title', 'API Documentation'),
-        'version': spec.get('info', {}).get('version', '1.0.0'),
-        'description': spec.get('info', {}).get('description', ''),
-        'endpoints': endpoints,
-        'tags': spec.get('tags', []),
-        'servers': spec.get('servers', [])
+def generate_template_context(spec: OpenAPISpec) -> dict[str, Any]:
+    spec_dict = spec.model_dump()
+    return {
+        "spec": spec_dict,
+        "info": spec_dict.get("info", {}),
+        "paths": spec_dict.get("paths", {}),
+        "servers": spec_dict.get("servers", []),
+        "tags": spec_dict.get("tags", []),
+        "components": spec_dict.get("components", {}),
+        "security": spec_dict.get("security", []),
+        "external_docs": spec_dict.get("externalDocs"),
     }
-
-    with open(output, 'w') as f:
-        json.dump(docs, f, indent=2)