"""
Phase 36: Multi-Step Workflow Engine — Full DAG Orchestration.

Ties together task_dependencies (Phase 15) and blueprints (Phase 25) into a
complete workflow execution engine. Supports workflow definition with DAG
validation, step-by-step execution with dependency tracking, templates,
history, and parallel execution groups.

Unlike the simpler Phase 35 workflow.py (which stores runs separately),
this engine keeps each workflow's definition + runtime state in a single file,
logs all events to a shared history JSONL, and supports templates for reuse.

Storage:
  state/broker/workflows/{workflow_id}.json       — workflow definition + state
  state/broker/workflows/templates/{name}.json    — reusable templates
  state/broker/workflows/history.jsonl            — event log
"""

import json
import uuid
from collections import deque
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Optional

from brynq.server import (
    BROKER_DIR,
    _write_json_sync,
    _read_jsonl_sync,
    _append_jsonl_sync,
)

# ── Paths ──────────────────────────────────────────────────────────────────

WORKFLOWS_DIR = BROKER_DIR / "workflows"
TEMPLATES_DIR = WORKFLOWS_DIR / "templates"
HISTORY_FILE = WORKFLOWS_DIR / "history.jsonl"

for _d in (WORKFLOWS_DIR, TEMPLATES_DIR):
    _d.mkdir(parents=True, exist_ok=True)

# ── Constants ──────────────────────────────────────────────────────────────

WORKFLOW_STATES = ("draft", "running", "completed", "failed", "cancelled")
STEP_STATES = ("pending", "assigned", "running", "completed", "failed", "skipped")

DEFAULT_TIMEOUT_MINUTES = 30


# ── DAG Validation ─────────────────────────────────────────────────────────


def _validate_dag(steps: list[dict]) -> Optional[str]:
    """Validate that the steps form a valid DAG (no cycles).

    Uses Kahn's algorithm for topological sort. Returns None if valid,
    or an error message string if a cycle is detected or references are invalid.
    """
    step_ids = {s["id"] for s in steps}

    # Check for references to non-existent steps
    for step in steps:
        for dep in step.get("depends_on", []):
            if dep not in step_ids:
                return f"Step '{step['id']}' depends on unknown step '{dep}'"

    # Build adjacency for Kahn's algorithm
    in_degree: dict[str, int] = {s["id"]: 0 for s in steps}
    dependents: dict[str, list[str]] = {s["id"]: [] for s in steps}

    for step in steps:
        for dep in step.get("depends_on", []):
            in_degree[step["id"]] += 1
            dependents[dep].append(step["id"])

    queue = deque(sid for sid, deg in in_degree.items() if deg == 0)
    sorted_count = 0

    while queue:
        current = queue.popleft()
        sorted_count += 1
        for dependent in dependents[current]:
            in_degree[dependent] -= 1
            if in_degree[dependent] == 0:
                queue.append(dependent)

    if sorted_count != len(steps):
        return "Cycle detected in workflow DAG"

    return None


# ── History Logging ────────────────────────────────────────────────────────


def _log_event(workflow_id: str, event_type: str, details: Optional[dict] = None) -> None:
    """Append an event to the workflow history log."""
    event = {
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "workflow_id": workflow_id,
        "event": event_type,
        "details": details or {},
    }
    _append_jsonl_sync(HISTORY_FILE, event)


# ── Internal Helpers ───────────────────────────────────────────────────────


def _load_workflow(workflow_id: str) -> Optional[dict]:
    """Load a workflow from disk."""
    path = WORKFLOWS_DIR / f"{workflow_id}.json"
    if not path.exists():
        return None
    try:
        return json.loads(path.read_text("utf-8"))
    except (json.JSONDecodeError, OSError):
        return None


def _save_workflow(workflow: dict) -> None:
    """Persist a workflow to disk."""
    wid = workflow["workflow_id"]
    _write_json_sync(WORKFLOWS_DIR / f"{wid}.json", workflow)


def _now() -> str:
    return datetime.now(timezone.utc).isoformat()


# ── Workflow Definition ────────────────────────────────────────────────────


def create_workflow(
    name: str,
    description: str = "",
    steps: Optional[list[dict]] = None,
) -> dict:
    """Create a new workflow definition.

    Each step dict should contain:
      - id:              unique step identifier
      - name:            human-readable step name
      - required_skills: list of skills needed (default: [])
      - depends_on:      list of step IDs that must complete first (default: [])
      - timeout_minutes: max time before auto-fail (default: 30)
      - config:          arbitrary dict passed to the executing agent (default: {})

    Returns the workflow dict on success, or {"error": "..."} on failure.
    """
    if not name:
        return {"error": "Workflow name is required"}

    steps = steps or []

    # Validate step structure
    step_ids: set[str] = set()
    validated: list[dict] = []

    for step in steps:
        sid = step.get("id", "")
        if not sid:
            return {"error": "Every step must have an 'id'"}
        if sid in step_ids:
            return {"error": f"Duplicate step ID: '{sid}'"}
        step_ids.add(sid)

        validated.append({
            "id": sid,
            "name": step.get("name", sid),
            "required_skills": step.get("required_skills", []),
            "depends_on": step.get("depends_on", []),
            "timeout_minutes": step.get("timeout_minutes", DEFAULT_TIMEOUT_MINUTES),
            "config": step.get("config", {}),
            "state": "pending",
            "assigned_to": None,
            "started_at": None,
            "completed_at": None,
            "result": None,
            "error": None,
        })

    # DAG validation
    if validated:
        dag_error = _validate_dag(validated)
        if dag_error:
            return {"error": dag_error}

    workflow_id = uuid.uuid4().hex[:12]
    now = _now()

    workflow = {
        "workflow_id": workflow_id,
        "name": name,
        "description": description,
        "steps": validated,
        "state": "draft",
        "created_at": now,
        "updated_at": now,
        "started_at": None,
        "completed_at": None,
    }

    _save_workflow(workflow)
    _log_event(workflow_id, "created", {"name": name, "step_count": len(validated)})

    return workflow


# ── Workflow Execution ─────────────────────────────────────────────────────


def start_workflow(workflow_id: str) -> dict:
    """Begin execution of a workflow.

    Transitions workflow from 'draft' to 'running'. Steps whose dependencies
    are already satisfied are marked 'pending' (ready for assignment).
    """
    workflow = _load_workflow(workflow_id)
    if not workflow:
        return {"error": f"Workflow '{workflow_id}' not found"}

    if workflow["state"] != "draft":
        return {"error": f"Workflow is in state '{workflow['state']}', must be 'draft' to start"}

    now = _now()
    workflow["state"] = "running"
    workflow["started_at"] = now
    workflow["updated_at"] = now

    # Mark steps with no dependencies as pending (ready to be picked up)
    for step in workflow["steps"]:
        if not step["depends_on"]:
            step["state"] = "pending"

    _save_workflow(workflow)
    _log_event(workflow_id, "started")

    return workflow


def get_ready_steps(workflow_id: str) -> list[dict]:
    """Return steps whose dependencies are all completed (or skipped).

    These steps are in 'pending' state and ready for assignment.
    """
    workflow = _load_workflow(workflow_id)
    if not workflow:
        return []

    step_map = {s["id"]: s for s in workflow["steps"]}
    ready = []

    for step in workflow["steps"]:
        if step["state"] != "pending":
            continue
        deps_met = all(
            step_map[dep]["state"] in ("completed", "skipped")
            for dep in step["depends_on"]
            if dep in step_map
        )
        if deps_met:
            ready.append(step)

    return ready


def complete_step(
    workflow_id: str,
    step_id: str,
    result: Any = None,
    success: bool = True,
) -> dict:
    """Mark a step as completed and unlock downstream steps.

    If success is False, delegates to fail_step instead.
    """
    if not success:
        return fail_step(workflow_id, step_id, str(result) if result else "Step failed")

    workflow = _load_workflow(workflow_id)
    if not workflow:
        return {"error": f"Workflow '{workflow_id}' not found"}

    step_map = {s["id"]: s for s in workflow["steps"]}
    step = step_map.get(step_id)
    if not step:
        return {"error": f"Step '{step_id}' not found in workflow '{workflow_id}'"}

    if step["state"] in ("completed", "failed", "skipped"):
        return {"error": f"Step '{step_id}' is already in terminal state '{step['state']}'"}

    now = _now()
    step["state"] = "completed"
    step["completed_at"] = now
    step["result"] = result

    # Unlock downstream steps: check all steps that depend on this one
    newly_ready = []
    for s in workflow["steps"]:
        if s["state"] != "pending":
            continue
        if step_id not in s["depends_on"]:
            continue
        # Check if ALL dependencies are now completed/skipped
        all_met = all(
            step_map[dep]["state"] in ("completed", "skipped")
            for dep in s["depends_on"]
            if dep in step_map
        )
        if all_met:
            newly_ready.append(s["id"])

    # Check if the entire workflow is done
    all_terminal = all(
        s["state"] in ("completed", "failed", "skipped")
        for s in workflow["steps"]
    )
    if all_terminal:
        workflow["state"] = "completed"
        workflow["completed_at"] = now

    workflow["updated_at"] = now
    _save_workflow(workflow)

    _log_event(workflow_id, "step_completed", {"step_id": step_id, "newly_ready": newly_ready})

    return {
        "workflow_id": workflow_id,
        "step_id": step_id,
        "step_state": "completed",
        "workflow_state": workflow["state"],
        "newly_ready": newly_ready,
        "progress": _calc_progress(workflow),
    }


def fail_step(
    workflow_id: str,
    step_id: str,
    error: str,
    cascade: bool = True,
) -> dict:
    """Mark a step as failed.

    If cascade is True (default), all steps that transitively depend on this
    step are marked 'skipped' and the workflow is marked 'failed'.
    If cascade is False, only this step fails; other paths may continue.
    """
    workflow = _load_workflow(workflow_id)
    if not workflow:
        return {"error": f"Workflow '{workflow_id}' not found"}

    step_map = {s["id"]: s for s in workflow["steps"]}
    step = step_map.get(step_id)
    if not step:
        return {"error": f"Step '{step_id}' not found in workflow '{workflow_id}'"}

    if step["state"] in ("completed", "failed", "skipped"):
        return {"error": f"Step '{step_id}' is already in terminal state '{step['state']}'"}

    now = _now()
    step["state"] = "failed"
    step["error"] = error
    step["completed_at"] = now

    skipped_steps: list[str] = []

    if cascade:
        # BFS to find all transitively dependent steps
        dependents_map: dict[str, list[str]] = {s["id"]: [] for s in workflow["steps"]}
        for s in workflow["steps"]:
            for dep in s["depends_on"]:
                if dep in dependents_map:
                    dependents_map[dep].append(s["id"])

        queue = deque(dependents_map[step_id])
        visited: set[str] = set()
        while queue:
            current = queue.popleft()
            if current in visited:
                continue
            visited.add(current)
            cs = step_map[current]
            if cs["state"] not in ("completed", "failed", "skipped"):
                cs["state"] = "skipped"
                cs["error"] = f"Skipped: upstream step '{step_id}' failed"
                skipped_steps.append(current)
            for d in dependents_map.get(current, []):
                if d not in visited:
                    queue.append(d)

        workflow["state"] = "failed"
        workflow["completed_at"] = now
    else:
        # Check if all non-failed paths are done
        all_terminal = all(
            s["state"] in ("completed", "failed", "skipped")
            for s in workflow["steps"]
        )
        if all_terminal:
            # If any step completed, mark as failed (partial); otherwise failed
            workflow["state"] = "failed"
            workflow["completed_at"] = now

    workflow["updated_at"] = now
    _save_workflow(workflow)

    _log_event(workflow_id, "step_failed", {
        "step_id": step_id,
        "error": error,
        "cascade": cascade,
        "skipped": skipped_steps,
    })

    return {
        "workflow_id": workflow_id,
        "step_id": step_id,
        "step_state": "failed",
        "workflow_state": workflow["state"],
        "skipped_steps": skipped_steps,
        "progress": _calc_progress(workflow),
    }


def cancel_workflow(workflow_id: str) -> dict:
    """Cancel a running workflow."""
    workflow = _load_workflow(workflow_id)
    if not workflow:
        return {"error": f"Workflow '{workflow_id}' not found"}

    if workflow["state"] not in ("draft", "running"):
        return {"error": f"Cannot cancel workflow in state '{workflow['state']}'"}

    now = _now()
    workflow["state"] = "cancelled"
    workflow["completed_at"] = now
    workflow["updated_at"] = now

    _save_workflow(workflow)
    _log_event(workflow_id, "cancelled")

    return {"workflow_id": workflow_id, "state": "cancelled"}


# ── Status & Progress ─────────────────────────────────────────────────────


def _calc_progress(workflow: dict) -> dict:
    """Calculate progress metrics for a workflow."""
    steps = workflow.get("steps", [])
    total = len(steps)
    if total == 0:
        return {"percent": 100.0, "completed": 0, "total": 0, "by_state": {}}

    by_state: dict[str, int] = {}
    for s in steps:
        state = s["state"]
        by_state[state] = by_state.get(state, 0) + 1

    completed = by_state.get("completed", 0)
    skipped = by_state.get("skipped", 0)
    failed = by_state.get("failed", 0)
    done = completed + skipped + failed
    percent = round((done / total) * 100, 1)

    return {
        "percent": percent,
        "completed": completed,
        "total": total,
        "by_state": by_state,
    }


def get_workflow_status(workflow_id: str) -> dict:
    """Return full status of a workflow: each step's state, progress, etc."""
    workflow = _load_workflow(workflow_id)
    if not workflow:
        return {"error": f"Workflow '{workflow_id}' not found"}

    progress = _calc_progress(workflow)

    step_details = []
    for s in workflow["steps"]:
        step_details.append({
            "id": s["id"],
            "name": s["name"],
            "state": s["state"],
            "depends_on": s["depends_on"],
            "started_at": s.get("started_at"),
            "completed_at": s.get("completed_at"),
            "result": s.get("result"),
            "error": s.get("error"),
        })

    return {
        "workflow_id": workflow_id,
        "name": workflow["name"],
        "state": workflow["state"],
        "progress": progress,
        "steps": step_details,
        "created_at": workflow["created_at"],
        "started_at": workflow.get("started_at"),
        "completed_at": workflow.get("completed_at"),
    }


# ── Execution Order ───────────────────────────────────────────────────────


def get_execution_order(workflow_id: str) -> dict:
    """Return topologically sorted step list with parallel execution groups.

    Each group is a list of step IDs that can run simultaneously.
    Groups are ordered: group N must complete before group N+1 can start.
    """
    workflow = _load_workflow(workflow_id)
    if not workflow:
        return {"error": f"Workflow '{workflow_id}' not found"}

    steps = workflow["steps"]
    if not steps:
        return {"workflow_id": workflow_id, "groups": [], "flat_order": []}

    step_map = {s["id"]: s for s in steps}
    in_degree: dict[str, int] = {s["id"]: 0 for s in steps}
    dependents: dict[str, list[str]] = {s["id"]: [] for s in steps}

    for s in steps:
        for dep in s["depends_on"]:
            if dep in in_degree:
                in_degree[s["id"]] += 1
                dependents[dep].append(s["id"])

    # Modified Kahn's: process layer-by-layer for parallel groups
    groups: list[list[str]] = []
    flat_order: list[str] = []
    current_layer = [sid for sid, deg in in_degree.items() if deg == 0]

    while current_layer:
        current_layer.sort()  # deterministic ordering within a group
        groups.append(current_layer)
        flat_order.extend(current_layer)

        next_layer: list[str] = []
        for sid in current_layer:
            for dep_id in dependents[sid]:
                in_degree[dep_id] -= 1
                if in_degree[dep_id] == 0:
                    next_layer.append(dep_id)
        current_layer = next_layer

    return {
        "workflow_id": workflow_id,
        "groups": groups,
        "flat_order": flat_order,
    }


# ── Workflow Templates ─────────────────────────────────────────────────────


def save_as_template(workflow_id: str, template_name: str) -> dict:
    """Save a workflow as a reusable template.

    Strips runtime state (step results, timestamps) and keeps only the
    structural definition.
    """
    workflow = _load_workflow(workflow_id)
    if not workflow:
        return {"error": f"Workflow '{workflow_id}' not found"}

    if not template_name:
        return {"error": "Template name is required"}

    # Strip runtime state from steps
    clean_steps = []
    for s in workflow["steps"]:
        clean_steps.append({
            "id": s["id"],
            "name": s["name"],
            "required_skills": s.get("required_skills", []),
            "depends_on": s["depends_on"],
            "timeout_minutes": s.get("timeout_minutes", DEFAULT_TIMEOUT_MINUTES),
            "config": s.get("config", {}),
        })

    template = {
        "template_name": template_name,
        "source_workflow_id": workflow_id,
        "description": workflow.get("description", ""),
        "steps": clean_steps,
        "created_at": _now(),
    }

    _write_json_sync(TEMPLATES_DIR / f"{template_name}.json", template)
    _log_event(workflow_id, "saved_as_template", {"template_name": template_name})

    return template


def list_templates() -> list[dict]:
    """List all available workflow templates."""
    templates = []
    for f in sorted(TEMPLATES_DIR.glob("*.json")):
        try:
            templates.append(json.loads(f.read_text("utf-8")))
        except (json.JSONDecodeError, OSError):
            continue
    return templates


def create_from_template(template_name: str, overrides: Optional[dict] = None) -> dict:
    """Create a new workflow from a template.

    overrides can include:
      - name: override the workflow name
      - description: override the description
      - step_configs: dict of step_id -> config dict to merge
    """
    path = TEMPLATES_DIR / f"{template_name}.json"
    if not path.exists():
        return {"error": f"Template '{template_name}' not found"}

    try:
        template = json.loads(path.read_text("utf-8"))
    except (json.JSONDecodeError, OSError):
        return {"error": f"Failed to read template '{template_name}'"}

    overrides = overrides or {}
    name = overrides.get("name", f"From template: {template_name}")
    description = overrides.get("description", template.get("description", ""))

    # Apply step config overrides
    steps = template["steps"]
    step_configs = overrides.get("step_configs", {})
    for step in steps:
        if step["id"] in step_configs:
            step["config"] = {**step.get("config", {}), **step_configs[step["id"]]}

    return create_workflow(name=name, description=description, steps=steps)


# ── Workflow Listing & History ─────────────────────────────────────────────


def list_workflows(status: Optional[str] = None, limit: int = 20) -> list[dict]:
    """List all workflows, optionally filtered by state."""
    workflows = []
    for f in sorted(WORKFLOWS_DIR.glob("*.json"), reverse=True):
        if f.name == "history.jsonl":
            continue
        try:
            wf = json.loads(f.read_text("utf-8"))
        except (json.JSONDecodeError, OSError):
            continue
        if status and wf.get("state") != status:
            continue
        workflows.append({
            "workflow_id": wf["workflow_id"],
            "name": wf["name"],
            "state": wf["state"],
            "step_count": len(wf.get("steps", [])),
            "created_at": wf["created_at"],
            "progress": _calc_progress(wf),
        })
        if len(workflows) >= limit:
            break
    return workflows


def get_workflow_history(workflow_id: str) -> list[dict]:
    """Return timeline of all state changes for a workflow."""
    all_events = _read_jsonl_sync(HISTORY_FILE)
    return [e for e in all_events if e.get("workflow_id") == workflow_id]
