"""Declarative schema for workflow topologies.

Spec 08.7 deliverable.

WorkflowSchema is a serialisable Pydantic description of a workflow's
topology — its steps, edges, and config.  It can be serialised to/from
JSON or YAML and round-trips cleanly for step types whose configuration
is fully data-driven.

Step types and serialisability
-------------------------------
+---------------+------------------+-------------------------------------------+
| StepKind      | Schema support   | Notes                                     |
+---------------+------------------+-------------------------------------------+
| echo          | Full             | prefix/suffix/delay captured in config    |
| http          | Full             | no callable; inputs drive the request     |
| transform     | Topology only    | field_mapping serialisable; type schemas  |
|               |                  | stored as JSON Schema dicts               |
| function      | Topology only    | function_ref (dotted path) recorded;      |
|               |                  | not importable/callable in 08.7           |
| subworkflow   | Topology only    | nested WorkflowSchema planned for 08.8   |
| agent         | Topology only    | agent_ref (dotted path) recorded          |
+---------------+------------------+-------------------------------------------+

"Topology only" steps can be serialised *from* a live Workflow, but
the schema loader (08.8) cannot yet reconstruct their callables from the
schema alone.  The schema still records all topology fields so the graph
can be visualised, validated, and partially rebuilt (e.g. by injecting
callables at load time via a registry).

Code-bound fields that are never serialised
-------------------------------------------
- FunctionStep callables
- SubWorkflowStep workflow_factory, input_mapper, output_mapper
- WorkflowAgent input_mapper, output_mapper, history_selector

These remain Python-only.  Their schema representations carry enough
metadata (function_ref, agent_ref) to support future registry lookups.
"""

from __future__ import annotations

from enum import Enum
from typing import TYPE_CHECKING, Any

from pydantic import BaseModel, ConfigDict, Field

from agentbyte.workflow.core.models import (
    EdgeCondition,
    JoinFailureMode,
    WorkflowCompletionMode,
    WorkflowStateConflictPolicy,
    WorkflowStateMode,
)

if TYPE_CHECKING:
    from agentbyte.workflow.core.workflow import Workflow


# ---------------------------------------------------------------------------
# Step kind enum
# ---------------------------------------------------------------------------

class StepKind(str, Enum):
    """Known step implementation kinds."""

    ECHO = "echo"
    TRANSFORM = "transform"
    HTTP = "http"
    FUNCTION = "function"
    SUBWORKFLOW = "subworkflow"
    AGENT = "agent"
    UNKNOWN = "unknown"


# ---------------------------------------------------------------------------
# Edge schema
# ---------------------------------------------------------------------------

class EdgeConditionSchema(BaseModel):
    """Serialisable form of an EdgeCondition."""

    type: str = "always"
    field: str | None = None
    operator: str | None = None
    value: Any | None = None

    model_config = ConfigDict(extra="forbid")

    @classmethod
    def from_edge_condition(cls, condition: EdgeCondition) -> "EdgeConditionSchema":
        return cls(
            type=condition.type,
            field=condition.field,
            operator=condition.operator,
            value=condition.value,
        )

    def to_edge_condition(self) -> EdgeCondition:
        return EdgeCondition(
            type=self.type,
            field=self.field,
            operator=self.operator,
            value=self.value,
        )


class EdgeSchema(BaseModel):
    """Serialisable representation of a workflow edge."""

    from_step: str
    to_step: str
    condition: EdgeConditionSchema = Field(default_factory=EdgeConditionSchema)

    model_config = ConfigDict(extra="forbid")


# ---------------------------------------------------------------------------
# Step schema
# ---------------------------------------------------------------------------

class StepSchema(BaseModel):
    """Serialisable representation of a single workflow step.

    ``kind`` identifies the step implementation class.  ``config`` carries
    step-type-specific configuration (prefix/suffix for echo steps, url/method
    for http steps, etc.).  For code-bound step types (function, subworkflow,
    agent), ``config`` contains reference strings (``function_ref``,
    ``agent_ref``) that allow a future registry-based loader to resolve the
    callable.
    """

    id: str
    kind: StepKind = StepKind.UNKNOWN
    name: str
    description: str | None = None
    tags: list[str] = Field(default_factory=list)
    max_retries: int = 0
    timeout_seconds: float | None = None
    join_failure_mode: JoinFailureMode = JoinFailureMode.FAIL_FAST
    # Step-type-specific configuration.  Validated loosely here; the builder
    # validates more strictly when reconstructing concrete step instances.
    config: dict[str, Any] = Field(default_factory=dict)

    model_config = ConfigDict(extra="forbid")


# ---------------------------------------------------------------------------
# Workflow schema
# ---------------------------------------------------------------------------

class WorkflowSchema(BaseModel):
    """Serialisable description of a workflow topology.

    This is the top-level document produced by ``Workflow.to_schema()`` and
    consumed by the loader in Spec 08.8.

    Fields
    ------
    name
        Workflow identifier — matches ``WorkflowConfig.name``.
    schema_version
        Schema format version.  ``"1"`` for this release.
    description
        Optional human-readable description.
    state_mode / state_conflict_policy
        Match ``WorkflowConfig`` fields.
    completion_mode
        Match ``WorkflowConfig.completion_mode``.
    initial_state
        Seed values passed to ``WorkflowConfig.initial_state``.
    start_step / end_steps
        Topology anchors.
    steps / edges
        Full graph description.
    metadata
        Arbitrary pass-through fields from ``WorkflowConfig.metadata``.
    """

    name: str
    schema_version: str = "1"
    description: str | None = None
    state_mode: WorkflowStateMode = WorkflowStateMode.IMMEDIATE
    state_conflict_policy: WorkflowStateConflictPolicy = WorkflowStateConflictPolicy.FAIL
    completion_mode: WorkflowCompletionMode = WorkflowCompletionMode.ANY_END
    initial_state: dict[str, Any] = Field(default_factory=dict)
    start_step: str
    end_steps: list[str]
    steps: list[StepSchema]
    edges: list[EdgeSchema] = Field(default_factory=list)
    metadata: dict[str, Any] = Field(default_factory=dict)

    model_config = ConfigDict(extra="forbid")

    # ------------------------------------------------------------------
    # Serialisation helpers
    # ------------------------------------------------------------------

    def to_json(self, *, indent: int = 2) -> str:
        """Serialise to a JSON string."""
        return self.model_dump_json(indent=indent)

    def to_dict(self) -> dict[str, Any]:
        """Serialise to a JSON-safe dict."""
        return self.model_dump(mode="json")

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "WorkflowSchema":
        """Deserialise from a dict (e.g. parsed from JSON or YAML)."""
        return cls.model_validate(data)

    @classmethod
    def from_json(cls, json_str: str) -> "WorkflowSchema":
        """Deserialise from a JSON string."""
        return cls.model_validate_json(json_str)


# ---------------------------------------------------------------------------
# Workflow → schema extractor
# ---------------------------------------------------------------------------

def _detect_step_kind(step: Any) -> StepKind:
    """Infer StepKind from the concrete step class name."""
    cls_name = type(step).__name__
    mapping = {
        "EchoStep": StepKind.ECHO,
        "HttpStep": StepKind.HTTP,
        "TransformStep": StepKind.TRANSFORM,
        "FunctionStep": StepKind.FUNCTION,
        "SubWorkflowStep": StepKind.SUBWORKFLOW,
        "AgentStep": StepKind.AGENT,
    }
    return mapping.get(cls_name, StepKind.UNKNOWN)


def _extract_step_config(step: Any, kind: StepKind) -> dict[str, Any]:
    """Extract serialisable config fields from a concrete step instance."""
    if kind == StepKind.ECHO:
        return {
            "prefix": getattr(step, "prefix", ""),
            "suffix": getattr(step, "suffix", ""),
            "delay_seconds": getattr(step, "delay_seconds", 0.0),
        }
    if kind == StepKind.HTTP:
        # HttpStep has no stored config beyond metadata — inputs drive requests.
        return {}
    if kind == StepKind.TRANSFORM:
        return {
            "field_mapping": getattr(step, "mappings", {}),
            "input_type_schema": step.input_type.model_json_schema() if hasattr(step, "input_type") else {},
            "output_type_schema": step.output_type.model_json_schema() if hasattr(step, "output_type") else {},
        }
    if kind == StepKind.FUNCTION:
        fn = getattr(step, "func", None)
        if fn is not None:
            module = getattr(fn, "__module__", None)
            qualname = getattr(fn, "__qualname__", None)
            ref = f"{module}.{qualname}" if module and qualname else None
        else:
            ref = None
        return {
            "function_ref": ref,
            "input_type": step.input_type.__name__ if hasattr(step, "input_type") else None,
            "output_type": step.output_type.__name__ if hasattr(step, "output_type") else None,
        }
    if kind == StepKind.SUBWORKFLOW:
        return {
            "input_type": step.input_type.__name__ if hasattr(step, "input_type") else None,
            "output_type": step.output_type.__name__ if hasattr(step, "output_type") else None,
            # workflow_factory and mappers are code-bound; not serialisable.
        }
    if kind == StepKind.AGENT:
        agent = getattr(step, "agent", None)
        if agent is not None:
            cls = type(agent)
            ref = f"{cls.__module__}.{cls.__qualname__}"
        else:
            ref = None
        return {"agent_ref": ref}
    return {}


def workflow_to_schema(workflow: "Workflow") -> WorkflowSchema:
    """Extract a ``WorkflowSchema`` from a live ``Workflow`` instance.

    This is a pure serialisation path — no callables are invoked.  Code-bound
    fields (function callables, mapper callables) are captured as reference
    strings where possible, or omitted with a note in ``config``.
    """
    step_schemas: list[StepSchema] = []
    for step in workflow.steps.values():
        kind = _detect_step_kind(step)
        config = _extract_step_config(step, kind)
        step_schemas.append(
            StepSchema(
                id=step.step_id,
                kind=kind,
                name=step.metadata.name,
                description=step.metadata.description,
                tags=list(step.metadata.tags),
                max_retries=step.metadata.max_retries,
                timeout_seconds=step.metadata.timeout_seconds,
                join_failure_mode=step.metadata.join_failure_mode,
                config=config,
            )
        )

    edge_schemas: list[EdgeSchema] = [
        EdgeSchema(
            from_step=edge.from_step,
            to_step=edge.to_step,
            condition=EdgeConditionSchema.from_edge_condition(edge.condition),
        )
        for edge in workflow.edges
    ]

    cfg = workflow.config
    return WorkflowSchema(
        name=cfg.name,
        description=cfg.description,
        state_mode=cfg.state_mode,
        state_conflict_policy=cfg.state_conflict_policy,
        completion_mode=cfg.completion_mode,
        initial_state=dict(cfg.initial_state),
        start_step=workflow.start_step_id or "",
        end_steps=list(workflow.end_step_ids),
        steps=step_schemas,
        edges=edge_schemas,
        metadata=dict(cfg.metadata),
    )


__all__ = [
    "StepKind",
    "EdgeConditionSchema",
    "EdgeSchema",
    "StepSchema",
    "WorkflowSchema",
    "workflow_to_schema",
]
