"""Workflow schema loader and builder.

Spec 08.8 deliverable.
"""

from __future__ import annotations

import copy
from dataclasses import dataclass, field
from typing import Any, Callable

from pydantic import BaseModel, create_model

from agentbyte.workflow.core.models import StepMetadata
from agentbyte.workflow.core.workflow import Workflow, WorkflowConfig
from agentbyte.workflow.schema import StepKind, StepSchema, WorkflowSchema
from agentbyte.workflow.schema_utils import get_python_type_from_json_schema_type
from agentbyte.workflow.steps import EchoStep, HttpStep, TransformStep
from agentbyte.workflow.steps.step import BaseStep

StepBuilder = Callable[[StepSchema], BaseStep]


class WorkflowSchemaBuildError(ValueError):
    """Raised when a workflow schema cannot be rebuilt into a live workflow."""


@dataclass(slots=True)
class WorkflowStepRegistry:
    """Registry for rebuilding code-bound workflow steps.

    Data-driven steps such as ``EchoStep`` and ``HttpStep`` can be rebuilt
    directly from ``WorkflowSchema``. Code-bound steps require explicit
    injection through this registry.
    """

    by_step_id: dict[str, StepBuilder | BaseStep] = field(default_factory=dict)
    function_by_ref: dict[str, StepBuilder | BaseStep] = field(default_factory=dict)
    function_by_id: dict[str, StepBuilder | BaseStep] = field(default_factory=dict)
    subworkflow_by_id: dict[str, StepBuilder | BaseStep] = field(default_factory=dict)
    agent_by_ref: dict[str, StepBuilder | BaseStep] = field(default_factory=dict)
    agent_by_id: dict[str, StepBuilder | BaseStep] = field(default_factory=dict)
    transform_by_id: dict[str, StepBuilder | BaseStep] = field(default_factory=dict)
    unknown_by_id: dict[str, StepBuilder | BaseStep] = field(default_factory=dict)

    def resolve(self, step_schema: StepSchema) -> BaseStep | None:
        """Resolve a concrete step instance or step builder for a schema step."""
        if step_schema.id in self.by_step_id:
            return _materialize_step(self.by_step_id[step_schema.id], step_schema)

        if step_schema.kind == StepKind.FUNCTION:
            ref = step_schema.config.get("function_ref")
            if isinstance(ref, str) and ref in self.function_by_ref:
                return _materialize_step(self.function_by_ref[ref], step_schema)
            if step_schema.id in self.function_by_id:
                return _materialize_step(self.function_by_id[step_schema.id], step_schema)
        elif step_schema.kind == StepKind.SUBWORKFLOW and step_schema.id in self.subworkflow_by_id:
            return _materialize_step(self.subworkflow_by_id[step_schema.id], step_schema)
        elif step_schema.kind == StepKind.AGENT:
            ref = step_schema.config.get("agent_ref")
            if isinstance(ref, str) and ref in self.agent_by_ref:
                return _materialize_step(self.agent_by_ref[ref], step_schema)
            if step_schema.id in self.agent_by_id:
                return _materialize_step(self.agent_by_id[step_schema.id], step_schema)
        elif step_schema.kind == StepKind.TRANSFORM and step_schema.id in self.transform_by_id:
            return _materialize_step(self.transform_by_id[step_schema.id], step_schema)
        elif step_schema.kind == StepKind.UNKNOWN and step_schema.id in self.unknown_by_id:
            return _materialize_step(self.unknown_by_id[step_schema.id], step_schema)

        return None


def _materialize_step(
    builder: StepBuilder | BaseStep,
    step_schema: StepSchema,
) -> BaseStep:
    if isinstance(builder, BaseStep):
        step = copy.deepcopy(builder)
    else:
        step = builder(step_schema)

    if not isinstance(step, BaseStep):
        msg = (
            f"Step registry entry for '{step_schema.id}' did not return a BaseStep; "
            f"got {type(step)!r}"
        )
        raise WorkflowSchemaBuildError(msg)

    if step.step_id != step_schema.id:
        msg = (
            f"Registry step for schema step '{step_schema.id}' returned step_id "
            f"'{step.step_id}'. The IDs must match."
        )
        raise WorkflowSchemaBuildError(msg)
    return step


class WorkflowBuilder:
    """Build a live ``Workflow`` from ``WorkflowSchema``."""

    def __init__(self, *, registry: WorkflowStepRegistry | None = None) -> None:
        self._registry = registry or WorkflowStepRegistry()

    def from_schema(self, schema: WorkflowSchema) -> Workflow:
        config = WorkflowConfig(
            name=schema.name,
            description=schema.description,
            initial_state=dict(schema.initial_state),
            metadata=dict(schema.metadata),
            completion_mode=schema.completion_mode,
            state_mode=schema.state_mode,
            state_conflict_policy=schema.state_conflict_policy,
        )
        workflow = Workflow(config)

        for step_schema in schema.steps:
            workflow.add_step(self._build_step(step_schema))

        for edge_schema in schema.edges:
            workflow.add_edge(
                edge_schema.from_step,
                edge_schema.to_step,
                condition=edge_schema.condition.to_edge_condition(),
            )

        workflow.set_start_step(schema.start_step)
        for end_step in schema.end_steps:
            workflow.add_end_step(end_step)

        validation = workflow.validate_workflow()
        if not validation.is_valid:
            msg = (
                f"Workflow built from schema '{schema.name}' is invalid: "
                + "; ".join(validation.errors)
            )
            raise WorkflowSchemaBuildError(msg)
        return workflow

    def _build_step(self, step_schema: StepSchema) -> BaseStep:
        metadata = StepMetadata(
            name=step_schema.name,
            description=step_schema.description,
            tags=list(step_schema.tags),
            max_retries=step_schema.max_retries,
            timeout_seconds=step_schema.timeout_seconds,
            join_failure_mode=step_schema.join_failure_mode,
        )
        kind = step_schema.kind

        if kind == StepKind.ECHO:
            return EchoStep(
                step_schema.id,
                metadata,
                prefix=str(step_schema.config.get("prefix", "")),
                suffix=str(step_schema.config.get("suffix", "")),
                delay_seconds=float(step_schema.config.get("delay_seconds", 0.0)),
            )

        if kind == StepKind.HTTP:
            return HttpStep(step_schema.id, metadata)

        if kind == StepKind.TRANSFORM:
            resolved = self._registry.resolve(step_schema)
            if resolved is not None:
                return resolved
            field_mapping = step_schema.config.get("field_mapping")
            input_schema = step_schema.config.get("input_type_schema")
            output_schema = step_schema.config.get("output_type_schema")
            if (
                isinstance(field_mapping, dict)
                and isinstance(input_schema, dict)
                and isinstance(output_schema, dict)
            ):
                input_type = _build_model_from_json_schema(
                    f"{step_schema.id.title().replace('_', '')}Input",
                    input_schema,
                )
                output_type = _build_model_from_json_schema(
                    f"{step_schema.id.title().replace('_', '')}Output",
                    output_schema,
                )
                return TransformStep(
                    step_schema.id,
                    metadata,
                    input_type=input_type,
                    output_type=output_type,
                    mappings=field_mapping,
                )
            msg = (
                f"Transform step '{step_schema.id}' is missing serializable transform "
                "configuration and cannot be rebuilt."
            )
            raise WorkflowSchemaBuildError(msg)

        resolved = self._registry.resolve(step_schema)
        if resolved is not None:
            return resolved

        if kind == StepKind.FUNCTION:
            msg = (
                f"Function step '{step_schema.id}' cannot be rebuilt from schema alone. "
                "Provide a WorkflowStepRegistry entry keyed by function_ref or step id."
            )
            raise WorkflowSchemaBuildError(msg)
        if kind == StepKind.SUBWORKFLOW:
            msg = (
                f"SubWorkflowStep '{step_schema.id}' requires a registry entry because "
                "workflow factories and mapper callables are Python-only."
            )
            raise WorkflowSchemaBuildError(msg)
        if kind == StepKind.AGENT:
            msg = (
                f"Agent step '{step_schema.id}' requires a registry entry because "
                "agent instances are Python-only."
            )
            raise WorkflowSchemaBuildError(msg)
        msg = (
            f"Unknown step kind '{kind.value}' for step '{step_schema.id}'. "
            "Provide a registry entry keyed by step id."
        )
        raise WorkflowSchemaBuildError(msg)


class WorkflowLoader:
    """Load ``Workflow`` instances from serialised schema formats."""

    def __init__(self, *, registry: WorkflowStepRegistry | None = None) -> None:
        self._builder = WorkflowBuilder(registry=registry)

    def from_schema(self, schema: WorkflowSchema) -> Workflow:
        return self._builder.from_schema(schema)

    def from_dict(self, data: dict[str, Any]) -> Workflow:
        return self.from_schema(WorkflowSchema.from_dict(data))

    def from_json(self, json_str: str) -> Workflow:
        return self.from_schema(WorkflowSchema.from_json(json_str))

    def from_yaml(self, yaml_str: str) -> Workflow:
        try:
            import yaml
        except ImportError as exc:
            msg = "YAML loading requires PyYAML to be installed."
            raise WorkflowSchemaBuildError(msg) from exc
        data = yaml.safe_load(yaml_str)
        if not isinstance(data, dict):
            msg = "YAML workflow schema must deserialize to a mapping."
            raise WorkflowSchemaBuildError(msg)
        return self.from_dict(data)


def _build_model_from_json_schema(name: str, schema: dict[str, Any]) -> type[BaseModel]:
    properties = schema.get("properties", {})
    required = set(schema.get("required", []))
    fields: dict[str, tuple[type[Any], Any]] = {}

    for field_name, field_schema in properties.items():
        if not isinstance(field_schema, dict):
            field_type = Any
        else:
            json_type = field_schema.get("type")
            field_type = get_python_type_from_json_schema_type(
                json_type if isinstance(json_type, str) else None
            )
        default = ... if field_name in required else None
        fields[field_name] = (field_type, default)

    return create_model(name, **fields)


__all__ = [
    "WorkflowSchemaBuildError",
    "WorkflowStepRegistry",
    "WorkflowBuilder",
    "WorkflowLoader",
]
