{# Jinja2 template for generating typed Prefect deployment SDK #}
"""
Prefect SDK - Auto-generated typed client for workspace deployments.

Generated at: {{ metadata.generation_time }}
Prefect version: {{ metadata.prefect_version }}
{% if metadata.workspace_name -%}
Workspace: {{ metadata.workspace_name }}
{% endif -%}
{% if metadata.api_url -%}
API URL: {{ metadata.api_url }}
{% endif %}
This file was auto-generated by `prefect sdk generate`.
Do not edit manually - regenerate after deployment changes.
"""
from datetime import datetime
from typing import TYPE_CHECKING, Any, Iterable, Literal{% if deployments %}, cast{% endif %}{% if deployments | length > 1 %}, overload{% endif %}
{%- if work_pool_typeddicts %}

from typing_extensions import TypedDict
{%- endif %}

if TYPE_CHECKING:
    from prefect.client.schemas import FlowRun
    from prefect.futures import PrefectFlowRunFuture

{% if deployment_names -%}
# =============================================================================
# DEPLOYMENT NAME LITERAL
# =============================================================================
# All deployment names for autocomplete
DeploymentName = Literal[
{% for name in deployment_names %}
    {{ name | repr }},
{% endfor %}
]
{% else %}
# No deployments found
DeploymentName = Literal[""]
{% endif %}

{% if work_pool_typeddicts -%}
# =============================================================================
# WORK POOL JOB VARIABLES TYPEDDICTS
# =============================================================================
{% for wp in work_pool_typeddicts %}
class {{ wp.class_name }}(TypedDict, total=False):
    """Job variables for work pool: {{ wp.original_name }}"""
{% for field in wp.fields %}
    {{ field.name }}: {{ field.python_type }}
{% endfor %}
{% if not wp.fields %}
    pass
{% endif %}


{% endfor %}
{% endif %}
{% if deployments -%}
# =============================================================================
# DEPLOYMENT CLASSES
# =============================================================================
{% for dep in deployments %}
class {{ dep.class_name }}:
    """
    Deployment: {{ dep.full_name }}
{% if dep.work_pool_name %}
    Work Pool: {{ dep.work_pool_name }}
{% endif %}
{% if dep.description %}

    {{ dep.description | wordwrap(76) | indent(4) }}
{% endif %}
    """

    def __init__(self) -> None:
        self._options: dict[str, Any] = {}

    def _copy_with_options(self, _sdk_options: dict[str, Any]) -> "{{ dep.class_name }}":
        """Create a new instance with the given options dict."""
        _sdk_new = {{ dep.class_name }}()
        _sdk_new._options = _sdk_options
        return _sdk_new

    def with_options(
        self,
        *,
        tags: Iterable[str] | None = None,
        idempotency_key: str | None = None,
        work_queue_name: str | None = None,
        as_subflow: bool | None = None,
        scheduled_time: datetime | None = None,
        flow_run_name: str | None = None,
    ) -> "{{ dep.class_name }}":
        """Create a new deployment handle with updated run options.

        Returns a new instance with merged options (does not mutate self).
        This matches the behavior of Flow.with_options() and Task.with_options().

        Note: timeout and poll_interval are not included here - use the
        PrefectFlowRunFuture.result() method for waiting.
        """
        _sdk_options = self._options.copy()
        if tags is not None:
            _sdk_options["tags"] = tags
        if idempotency_key is not None:
            _sdk_options["idempotency_key"] = idempotency_key
        if work_queue_name is not None:
            _sdk_options["work_queue_name"] = work_queue_name
        if as_subflow is not None:
            _sdk_options["as_subflow"] = as_subflow
        if scheduled_time is not None:
            _sdk_options["scheduled_time"] = scheduled_time
        if flow_run_name is not None:
            _sdk_options["flow_run_name"] = flow_run_name
        return self._copy_with_options(_sdk_options)
{% if dep.has_job_variables %}

    def with_infra(
        self,
        *,
{% for field in dep.job_variable_fields %}
        {{ field.name }}: {{ field.python_type | make_optional }} = None,
{% endfor %}
    ) -> "{{ dep.class_name }}":
        """Create a new deployment handle with updated job variables.

        Returns a new instance with merged options (does not mutate self).
        Job variable types are derived from the work pool schema.
{%- if dep.job_variable_fields %}

        Args:
{%- for field in dep.job_variable_fields %}
            {{ field.name }}: {{ field.description if field.description else "(No description)" }}
{%- endfor %}
{%- endif %}
        """
        _sdk_options = self._options.copy()
        _sdk_job_vars = _sdk_options.get("job_variables", {}).copy()
{% for field in dep.job_variable_fields %}
        if {{ field.name }} is not None:
            _sdk_job_vars[{{ field.original_name | repr }}] = {{ field.name }}
{% endfor %}
        if _sdk_job_vars:
            _sdk_options["job_variables"] = _sdk_job_vars
        return self._copy_with_options(_sdk_options)
{% endif %}

    def run(
        self,
{% for param in dep.required_params %}
        {{ param.name }}: {{ param.python_type }},
{% endfor %}
{% for param in dep.optional_params %}
        {{ param.name }}: {% if param.has_default %}{{ param.python_type }} = {{ param.default_repr }}{% else %}{{ param.python_type | make_optional }} = None{% endif %},
{% endfor %}
    ) -> "PrefectFlowRunFuture[Any]":
        """Run the {{ dep.full_name }} deployment synchronously.
{%- if dep.required_params or dep.optional_params %}

        Args:
{%- for param in dep.required_params %}
            {{ param.name }}: {{ param.description if param.description else "(No description)" }}
{%- endfor %}
{%- for param in dep.optional_params %}
            {{ param.name }}: {{ param.description if param.description else "(No description)" }}
{%- endfor %}
{%- endif %}

        Returns:
            A PrefectFlowRunFuture that can be used to get the flow_run_id
            immediately, call .result() to wait for completion, or check .state.
        """
        from prefect.deployments import run_deployment
        from prefect.futures import PrefectFlowRunFuture

        _sdk_params: dict[str, Any] = {}
{% for param in dep.required_params %}
        _sdk_params[{{ param.original_name | repr }}] = {{ param.name }}
{% endfor %}
{% for param in dep.optional_params %}
{% if param.has_default %}
        if {{ param.name }} != {{ param.default_repr }}:
            _sdk_params[{{ param.original_name | repr }}] = {{ param.name }}
{% else %}
        if {{ param.name }} is not None:
            _sdk_params[{{ param.original_name | repr }}] = {{ param.name }}
{% endif %}
{% endfor %}

        # _sync=True ensures @async_dispatch returns sync result, not coroutine
        _sdk_flow_run = cast(
            "FlowRun",
            run_deployment(
                name={{ dep.full_name | repr }},
                parameters=_sdk_params,
                timeout=0,
                _sync=True,  # pyright: ignore[reportCallIssue]
                **self._options,
            ),
        )
        return PrefectFlowRunFuture(flow_run_id=_sdk_flow_run.id)

    async def run_async(
        self,
{% for param in dep.required_params %}
        {{ param.name }}: {{ param.python_type }},
{% endfor %}
{% for param in dep.optional_params %}
        {{ param.name }}: {% if param.has_default %}{{ param.python_type }} = {{ param.default_repr }}{% else %}{{ param.python_type | make_optional }} = None{% endif %},
{% endfor %}
    ) -> "PrefectFlowRunFuture[Any]":
        """Run the {{ dep.full_name }} deployment asynchronously.
{%- if dep.required_params or dep.optional_params %}

        Args:
{%- for param in dep.required_params %}
            {{ param.name }}: {{ param.description if param.description else "(No description)" }}
{%- endfor %}
{%- for param in dep.optional_params %}
            {{ param.name }}: {{ param.description if param.description else "(No description)" }}
{%- endfor %}
{%- endif %}

        Returns:
            A PrefectFlowRunFuture that can be used to get the flow_run_id
            immediately, call .result() to wait for completion, or check .state.
        """
        from prefect.deployments import arun_deployment
        from prefect.futures import PrefectFlowRunFuture

        _sdk_params: dict[str, Any] = {}
{% for param in dep.required_params %}
        _sdk_params[{{ param.original_name | repr }}] = {{ param.name }}
{% endfor %}
{% for param in dep.optional_params %}
{% if param.has_default %}
        if {{ param.name }} != {{ param.default_repr }}:
            _sdk_params[{{ param.original_name | repr }}] = {{ param.name }}
{% else %}
        if {{ param.name }} is not None:
            _sdk_params[{{ param.original_name | repr }}] = {{ param.name }}
{% endif %}
{% endfor %}

        _sdk_flow_run = await arun_deployment(
            name={{ dep.full_name | repr }},
            parameters=_sdk_params,
            timeout=0,
            **self._options,
        )
        return PrefectFlowRunFuture(flow_run_id=_sdk_flow_run.id)


{% endfor %}
{% endif %}
# =============================================================================
# DEPLOYMENT NAMESPACE
# =============================================================================
class deployments:
    """
    Access deployments by name.

    Usage:
        from {{ module_name }} import deployments

        future = deployments.from_name("flow/deployment").with_options(
            tags=["production"],
        ).run(
            param="value",
        )
        result = future.result()  # Wait for completion

{% if deployment_names %}
    Available deployments:
{% for name in deployment_names %}
        - {{ name }}
{% endfor %}
{% else %}
    No deployments available.
{% endif %}
    """
{% if deployments %}
{% if deployments | length > 1 %}
{% for dep in deployments %}

    @overload
    @staticmethod
    def from_name(name: Literal[{{ dep.full_name | repr }}]) -> {{ dep.class_name }}: ...
{% endfor %}

    @staticmethod
    def from_name(name: DeploymentName) -> {% for dep in deployments %}{{ dep.class_name }}{% if not loop.last %} | {% endif %}{% endfor %}:
{% else %}
    @staticmethod
    def from_name(name: Literal[{{ deployments[0].full_name | repr }}]) -> {{ deployments[0].class_name }}:
{% endif %}
        """Get a deployment by name.

        Args:
            name: The deployment name in "flow-name/deployment-name" format.

        Returns:
            A deployment object with run() and run_async() methods.

        Raises:
            KeyError: If the deployment name is not found.
        """
        _sdk_deployments: dict[str, Any] = {
{% for dep in deployments %}
            {{ dep.full_name | repr }}: {{ dep.class_name }}(),
{% endfor %}
        }
        return _sdk_deployments[name]
{% else %}
    @staticmethod
    def from_name(name: DeploymentName) -> None:
        """Get a deployment by name.

        No deployments are available.
        """
        raise KeyError(f"Deployment not found: {name}")
{% endif %}


__all__ = ["deployments", "DeploymentName"]
