"""Schema extractor for Modal sandbox execution.

This module extracts JSON schemas from @pipeline(Schema) decorated functions.
It runs inside Modal sandboxes where exec() is safe.

Usage from Modal:
    python -m alloy_runtime_sdk.pipeline.schema_extractor

Environment Variables:
    ALLOY_RUNTIME_MODULE_PATH: File path to the Python source code
    ALLOY_RUNTIME_PIPELINE_NAME: Name of the @pipeline function (optional)

Output (stdout):
    {"success": true, "schema": {...}}  # Schema extracted
    {"success": true, "schema": null}   # No schema in @pipeline decorator
    {"success": false, "error": "..."}  # Error occurred
"""

import json
import os
import sys
from pathlib import Path
from typing import Any

from alloy_runtime_sdk.pipeline.decorators import (
    clear_pipeline_registry,
    get_pipeline_registry,
)


class SchemaExtractorError(Exception):
    """Raised when schema extraction encounters an error."""

    pass


def extract_schema_from_module(
    module_content: str,
    pipeline_name: str | None,
) -> dict[str, Any] | None:
    """Extract JSON schema from a pipeline module.

    Executes the module content to register @pipeline functions, then
    extracts the JSON schema if a Pydantic model was passed to @pipeline.

    Args:
        module_content: Python source code containing @pipeline decorated functions.
        pipeline_name: Name of the specific pipeline to extract schema from.
                      If None, uses the first/only pipeline.

    Returns:
        JSON Schema dict if @pipeline(Schema) used, else None.

    Raises:
        SchemaExtractorError: If module cannot be loaded or pipeline not found.
    """
    # Clear any existing registrations
    clear_pipeline_registry()

    # Create a module namespace with necessary imports
    namespace: dict[str, Any] = {
        "__name__": "__pipeline_module__",
        "__builtins__": __builtins__,
    }

    try:
        # Execute the module code to register @pipeline functions
        exec(module_content, namespace)
    except SyntaxError as e:
        raise SchemaExtractorError(f"Syntax error in pipeline module: {e}") from e
    except Exception as e:
        raise SchemaExtractorError(f"Failed to load pipeline module: {e}") from e

    # Get registered pipelines
    registry = get_pipeline_registry()
    if not registry:
        raise SchemaExtractorError(
            "No @pipeline decorated functions found in module. "
            "Make sure your function is decorated with @pipeline."
        )

    # Find the target pipeline
    if pipeline_name:
        if pipeline_name not in registry:
            available = list(registry.keys())
            raise SchemaExtractorError(
                f"Pipeline '{pipeline_name}' not found. Available: {available}"
            )
        pipeline_func = registry[pipeline_name]
    else:
        # No name specified - use the first one
        pipeline_func = list(registry.values())[0]

    # Extract schema from _input_schema_class if present
    # The _input_schema_class attribute is set when @pipeline(Schema) is used
    schema_class = getattr(pipeline_func, "_input_schema_class", None)
    if schema_class is not None:
        # Use Pydantic's model_json_schema() to get JSON Schema
        return schema_class.model_json_schema()

    # No schema in @pipeline decorator
    return None


def main() -> None:
    """Main entry point for Modal sandbox execution.

    Reads module content from environment variable, extracts schema,
    and outputs result as JSON to stdout.
    """
    module_path_str = os.environ.get("ALLOY_RUNTIME_MODULE_PATH")
    if not module_path_str:
        print(
            json.dumps(
                {
                    "success": False,
                    "error": "ALLOY_RUNTIME_MODULE_PATH environment variable is required",
                }
            )
        )
        sys.exit(1)

    try:
        module_content = Path(module_path_str).read_text(encoding="utf-8")
    except OSError as e:
        print(
            json.dumps(
                {
                    "success": False,
                    "error": f"Failed to read module content: {e}",
                }
            )
        )
        sys.exit(1)

    # Optional pipeline name
    pipeline_name = os.environ.get("ALLOY_RUNTIME_PIPELINE_NAME")

    # Extract schema
    try:
        schema = extract_schema_from_module(module_content, pipeline_name)

        print(
            json.dumps(
                {
                    "success": True,
                    "schema": schema,
                }
            )
        )

    except SchemaExtractorError as e:
        print(
            json.dumps(
                {
                    "success": False,
                    "error": str(e),
                }
            )
        )
        sys.exit(1)

    except Exception as e:
        print(
            json.dumps(
                {
                    "success": False,
                    "error": f"Unexpected error: {e}",
                }
            )
        )
        sys.exit(1)


if __name__ == "__main__":
    main()
