"""
Workflow definition schema and helpers for project-scoped reusable workflows.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from datetime import datetime
import uuid
from typing import Any, Dict, List, Optional

import copy

from .models import EngineType, Provenance, WorkflowStep


class WorkflowDefinitionStepType:
    """Supported workflow step types for the first workflow-engine pass."""

    TASK = "task"
    SET_VARIABLE = "set_variable"
    FOR_LOOP = "for_loop"
    IF_ELSE = "if_else"
    LOG = "log"
    TRY_CATCH = "try_catch"
    # Raster-support Phase 4a: raster-algebra one-shot step. Per-band
    # and per-cell iteration are loop_type values inside FOR_LOOP
    # ("each_band" / "each_cell"), mirroring "each_feature".
    RASTER_ALGEBRA = "raster_algebra"

    ALL = [TASK, SET_VARIABLE, FOR_LOOP, IF_ELSE, LOG, TRY_CATCH, RASTER_ALGEBRA]

    # Step types exposed in Pipeline Builder insert menus (raster primitives hidden
    # until ArcGIS Pro executor parity — see UX spec G.2).
    UI_INSERT_TYPES = [TASK, SET_VARIABLE, FOR_LOOP, IF_ELSE, LOG, TRY_CATCH]

    # Raster-specific loop modes hidden from the loop-type picker.
    UI_HIDDEN_LOOP_TYPES = frozenset({"each_band", "each_cell"})


@dataclass
class WorkflowDefinitionStep:
    """Serializable workflow-definition step."""

    id: str = field(default_factory=lambda: str(uuid.uuid4()))
    name: str = ""
    description: str = ""
    step_type: str = WorkflowDefinitionStepType.TASK
    enabled: bool = True
    continue_on_error: bool = False
    config: Dict[str, Any] = field(default_factory=dict)
    children: List["WorkflowDefinitionStep"] = field(default_factory=list)
    else_children: List["WorkflowDefinitionStep"] = field(default_factory=list)

    def to_dict(self) -> Dict[str, Any]:
        return {
            "id": self.id,
            "name": self.name,
            "description": self.description,
            "step_type": self.step_type,
            "enabled": self.enabled,
            "continue_on_error": self.continue_on_error,
            "config": dict(self.config or {}),
            "children": [child.to_dict() for child in self.children],
            "else_children": [child.to_dict() for child in self.else_children],
        }

    @classmethod
    def from_dict(cls, payload: Optional[Dict[str, Any]]) -> "WorkflowDefinitionStep":
        payload = dict(payload or {})
        children = [
            cls.from_dict(child_payload)
            for child_payload in list(payload.get("children", []) or [])
        ]
        else_children = [
            cls.from_dict(child_payload)
            for child_payload in list(payload.get("else_children", []) or [])
        ]
        return cls(
            id=str(payload.get("id") or str(uuid.uuid4())),
            name=str(payload.get("name", "") or "").strip(),
            description=str(payload.get("description", "") or ""),
            step_type=str(payload.get("step_type") or WorkflowDefinitionStepType.TASK).strip(),
            enabled=bool(payload.get("enabled", True)),
            continue_on_error=bool(payload.get("continue_on_error", False)),
            config=dict(payload.get("config", {}) or {}),
            children=children,
            else_children=else_children,
        )


@dataclass
class WorkflowDefinition:
    """Serializable definition stored on provenance metadata."""

    id: str = field(default_factory=lambda: str(uuid.uuid4()))
    name: str = ""
    description: str = ""
    version: int = 1
    steps: List[WorkflowDefinitionStep] = field(default_factory=list)
    created_at: str = field(default_factory=lambda: datetime.now().isoformat())
    updated_at: str = field(default_factory=lambda: datetime.now().isoformat())

    def touch(self) -> None:
        self.updated_at = datetime.now().isoformat()

    def to_dict(self) -> Dict[str, Any]:
        self.touch()
        return {
            "id": self.id,
            "name": self.name,
            "description": self.description,
            "version": self.version,
            "created_at": self.created_at,
            "updated_at": self.updated_at,
            "steps": [step.to_dict() for step in self.steps],
        }

    @classmethod
    def from_dict(cls, payload: Optional[Dict[str, Any]]) -> "WorkflowDefinition":
        payload = dict(payload or {})
        steps = [
            WorkflowDefinitionStep.from_dict(step_payload)
            for step_payload in list(payload.get("steps", []) or [])
        ]
        return cls(
            id=str(payload.get("id") or str(uuid.uuid4())),
            name=str(payload.get("name", "") or "").strip(),
            description=str(payload.get("description", "") or "").strip(),
            version=int(payload.get("version", 1) or 1),
            steps=steps,
            created_at=str(payload.get("created_at") or datetime.now().isoformat()),
            updated_at=str(payload.get("updated_at") or datetime.now().isoformat()),
        )


_ENGINE_TO_PIPELINE_KEY = {
    EngineType.NETWEAVER.value: "netweaver",
    EngineType.CDP.value: "cdp",
    EngineType.TEMPLATE_1.value: "lpa",
    EngineType.TEMPLATE_2.value: "portfolio",
    EngineType.WORKFLOW.value: "workflow",
    EngineType.PYTHON_SCRIPT.value: "python_script",
}


def _task_config_from_workflow_step(workflow_step: WorkflowStep, *, is_root: bool) -> Dict[str, Any]:
    """Minimal TASK config synthesized from a tree task when no pipeline step exists."""
    params = dict(workflow_step.parameters or {})
    engine_value = (
        workflow_step.engine_type.value
        if hasattr(workflow_step.engine_type, "value")
        else str(workflow_step.engine_type)
    )
    engine_key = _ENGINE_TO_PIPELINE_KEY.get(engine_value, str(engine_value).strip().lower())
    config: Dict[str, Any] = {
        "task_id": workflow_step.id,
        "engine_type": engine_key,
        "model_reference": str(params.get("model_file") or ""),
        "input_source": "assessment" if is_root else "previous_result",
    }
    if params.get("field_mappings") or params.get("topics") or params.get("configured_via_wizard"):
        config["configured_via_wizard"] = True
        config["task_parameters"] = copy.deepcopy(params)
    return config


def _build_task_step_from_workflow_step(
    workflow_step: WorkflowStep,
    existing_by_task_id: Optional[Dict[str, WorkflowDefinitionStep]] = None,
    *,
    is_root: bool = True,
) -> WorkflowDefinitionStep:
    """Build a pipeline TASK step from a tree task, preserving rich config when known."""
    existing = (existing_by_task_id or {}).get(workflow_step.id)
    if existing is not None and existing.step_type == WorkflowDefinitionStepType.TASK:
        step = WorkflowDefinitionStep.from_dict(existing.to_dict())
        step.name = workflow_step.name or step.name
        step.config = dict(step.config or {})
        step.config["task_id"] = workflow_step.id
    else:
        step = WorkflowDefinitionStep(
            name=workflow_step.name,
            step_type=WorkflowDefinitionStepType.TASK,
            config=_task_config_from_workflow_step(workflow_step, is_root=is_root),
        )
    step.children = [
        _build_task_step_from_workflow_step(child_step, existing_by_task_id, is_root=False)
        for child_step in list(workflow_step.children or [])
    ]
    step.else_children = []
    return step


def _index_tree_tasks_by_id(provenance: Provenance) -> Dict[str, WorkflowStep]:
    """Map every tree task id (recursive) to its WorkflowStep."""
    index: Dict[str, WorkflowStep] = {}

    def walk(workflow_step: WorkflowStep) -> None:
        if getattr(workflow_step, "id", None):
            index[str(workflow_step.id)] = workflow_step
        for child in list(getattr(workflow_step, "children", []) or []):
            walk(child)

    for root in list(provenance.workflow_steps or []):
        walk(root)
    return index


def sync_pipeline_from_task_tree(
    provenance: Provenance,
    existing_definition: WorkflowDefinition,
) -> WorkflowDefinition:
    """Level A: reconcile the linked pipeline with the task tree, non-destructively.

    The task tree is the inventory of engine TASK nodes; the pipeline is the
    authoritative *structure* and may contain user-authored content the tree
    cannot express. This reconcile therefore PRESERVES:

    * control-flow steps (FOR_LOOP / IF_ELSE / LOG / TRY_CATCH / SET_VARIABLE /
      RASTER_ALGEBRA) and everything nested inside them;
    * scenario-variant TASK nodes materialized from the Scenario Builder
      (``config['scenario_variant']``), which have no baseline tree task;
    * the structure/ordering of TASK steps the user moved (e.g. into a loop).

    It only DROPS a top-level plain TASK step whose ``task_id`` no longer exists
    in the tree (a genuine task deletion), and APPENDS tree tasks that are not
    yet represented anywhere in the pipeline (newly added tasks). Existing TASK
    config is preserved by ``task_id``; the display name is refreshed from the
    tree. This is what lets hand-authored control flow survive a task save —
    fixing the prior rebuild-from-scratch data-loss footgun.
    """
    tree_tasks_by_id = _index_tree_tasks_by_id(provenance)
    seen_task_ids: set = set()
    # task_id -> the reconciled pipeline step that represents it, so newly-added
    # tree children can be attached UNDER their parent (wherever it lives, e.g.
    # inside a loop) instead of being promoted to the top level.
    reconciled_by_task_id: Dict[str, WorkflowDefinitionStep] = {}

    def reconcile(
        step: WorkflowDefinitionStep,
        *,
        top_level: bool,
    ) -> Optional[WorkflowDefinitionStep]:
        if step.step_type == WorkflowDefinitionStepType.TASK:
            task_id = str((step.config or {}).get("task_id") or "").strip()
            tree_task = tree_tasks_by_id.get(task_id) if task_id else None
            is_scenario_variant = bool((step.config or {}).get("scenario_variant"))

            if tree_task is None:
                # No baseline tree task. Preserve user-authored content:
                # scenario variants, and anything the user nested inside
                # control flow. Only a top-level plain task that lost its
                # tree node is treated as a deletion and dropped.
                if not (is_scenario_variant or not top_level):
                    return None
                preserved = WorkflowDefinitionStep.from_dict(step.to_dict())
                preserved.children = _reconcile_children(step.children)
                preserved.else_children = _reconcile_children(step.else_children)
                return preserved

            # Matching tree task: keep the step (preserving rich config) and
            # refresh its display name from the tree.
            seen_task_ids.add(task_id)
            updated = WorkflowDefinitionStep.from_dict(step.to_dict())
            updated.name = tree_task.name or updated.name
            updated.config = dict(updated.config or {})
            updated.config["task_id"] = task_id
            updated.children = _reconcile_children(step.children)
            updated.else_children = _reconcile_children(step.else_children)
            reconciled_by_task_id[task_id] = updated
            return updated

        # Control-flow step: always preserved, children reconciled.
        preserved = WorkflowDefinitionStep.from_dict(step.to_dict())
        preserved.children = _reconcile_children(step.children)
        preserved.else_children = _reconcile_children(step.else_children)
        return preserved

    def _reconcile_children(
        steps: List[WorkflowDefinitionStep],
    ) -> List[WorkflowDefinitionStep]:
        reconciled: List[WorkflowDefinitionStep] = []
        for child in list(steps or []):
            result = reconcile(child, top_level=False)
            if result is not None:
                reconciled.append(result)
        return reconciled

    steps: List[WorkflowDefinitionStep] = []
    for root in list(existing_definition.steps or []):
        result = reconcile(root, top_level=True)
        if result is not None:
            steps.append(result)

    # Append tree tasks not yet anywhere in the pipeline (newly added tasks).
    # A new task is nested UNDER its parent's reconciled step when that parent is
    # represented (e.g. the parent was moved into a loop), so a child is never
    # silently promoted to a top-level sibling. Pre-order so a new parent is
    # placed before its new children.
    def attach_new(workflow_step: WorkflowStep, parent_task_id: Optional[str]) -> None:
        sid = str(workflow_step.id)
        if sid in seen_task_ids:
            # Already represented; surface any unseen descendants under it.
            for child in list(getattr(workflow_step, "children", []) or []):
                attach_new(child, sid)
            return
        seen_task_ids.add(sid)
        new_step = WorkflowDefinitionStep(
            name=workflow_step.name,
            step_type=WorkflowDefinitionStepType.TASK,
            config=_task_config_from_workflow_step(workflow_step, is_root=parent_task_id is None),
        )
        parent_step = reconciled_by_task_id.get(parent_task_id) if parent_task_id else None
        if parent_step is not None:
            parent_step.children.append(new_step)
        else:
            steps.append(new_step)
        reconciled_by_task_id[sid] = new_step
        for child in list(getattr(workflow_step, "children", []) or []):
            attach_new(child, sid)

    for workflow_step in list(provenance.workflow_steps or []):
        attach_new(workflow_step, None)

    return WorkflowDefinition(
        id=existing_definition.id,
        name=existing_definition.name or f"{provenance.name} Pipeline",
        description=existing_definition.description or provenance.description,
        version=existing_definition.version,
        steps=steps,
        created_at=existing_definition.created_at,
        updated_at=existing_definition.updated_at,
    )


def build_scenario_task_step(
    working_step: WorkflowStep,
    *,
    source_task_id: str = "",
    scenario_name: str = "",
    is_root: bool = True,
) -> WorkflowDefinitionStep:
    """Materialize a Scenario Builder working task as a self-contained pipeline TASK node.

    The node carries the edited model/parameters via the wizard fast-path
    (``configured_via_wizard`` + ``task_parameters``) so it is directly runnable
    by the executor without a baseline tree task, and is flagged
    ``scenario_variant`` so :func:`sync_pipeline_from_task_tree` preserves it.
    """
    config = _task_config_from_workflow_step(working_step, is_root=is_root)
    config["configured_via_wizard"] = True
    config["task_parameters"] = copy.deepcopy(dict(working_step.parameters or {}))
    config["scenario_variant"] = True
    config["task_id"] = str(getattr(working_step, "id", "") or uuid.uuid4())
    if source_task_id:
        config["scenario_baseline_task_id"] = str(source_task_id)
    return WorkflowDefinitionStep(
        name=scenario_name or working_step.name,
        step_type=WorkflowDefinitionStepType.TASK,
        config=config,
    )


def build_default_workflow_definition(provenance: Provenance) -> WorkflowDefinition:
    """Create a first-pass definition from the existing provenance task structure."""

    definition = WorkflowDefinition(
        name=provenance.name,
        description=provenance.description,
        steps=[
            _build_task_step_from_workflow_step(workflow_step, is_root=True)
            for workflow_step in list(provenance.workflow_steps or [])
        ],
    )
    return definition


def load_workflow_definition_from_provenance(provenance: Provenance) -> WorkflowDefinition:
    """Load persisted definition or synthesize a default one from current tasks."""

    metadata = dict(getattr(provenance, "metadata", {}) or {})
    payload = metadata.get("workflow_definition")
    if isinstance(payload, dict):
        definition = WorkflowDefinition.from_dict(payload)
        if not definition.name:
            definition.name = provenance.name
        if not definition.description:
            definition.description = provenance.description
        return definition

    return build_default_workflow_definition(provenance)


def save_workflow_definition_to_provenance(
    provenance: Provenance,
    workflow_definition: WorkflowDefinition,
) -> None:
    """Legacy inline persistence kept for backward-compatibility reads."""

    metadata = dict(getattr(provenance, "metadata", {}) or {})
    metadata["workflow_definition"] = workflow_definition.to_dict()
    provenance.metadata = metadata
