"""Workflow Store — DAG-based workflow engine for the Brynq runtime.

Ported from src/brynq/workflow.py. Stores workflow definitions and runs
under ~/.brynq/workflows/ so the runtime operates independently of the
legacy broker.

Features:
  - DEFINE: Create workflows with steps, dependencies, and conditions
  - EXECUTE: Start a workflow run; steps with no deps are immediately ready
  - BRANCH: Conditional steps based on previous step results
  - PARALLEL: Fan-out steps that run concurrently
  - MONITOR: Track workflow progress, step statuses, timing

Storage:
  ~/.brynq/workflows/
    definitions/{workflow_id}.json  — workflow templates
    runs/{run_id}.json             — active/completed workflow runs
"""

import json
import uuid
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional

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

WORKFLOWS_DIR = Path.home() / ".brynq" / "workflows"
DEFINITIONS_DIR = WORKFLOWS_DIR / "definitions"
RUNS_DIR = WORKFLOWS_DIR / "runs"

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

VALID_STEP_TYPES = ("task", "condition", "parallel", "notify", "wait")
VALID_RUN_STATUSES = ("pending", "running", "completed", "failed", "cancelled")
VALID_STEP_STATUSES = ("pending", "ready", "running", "completed", "failed", "skipped")


# ── Helpers ────────────────────────────────────────────────────────────────


def _write_json(path: Path, data: dict) -> None:
    """Write a dict as formatted JSON to *path*."""
    path.write_text(json.dumps(data, indent=2, default=str), encoding="utf-8")


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


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

    Each step dict should have:
      - id: unique step identifier within the workflow
      - type: task, condition, parallel, notify, wait
      - title: human-readable step name
      - depends_on: list of step IDs that must complete first (default: [])
      - config: step-type-specific configuration
      - inputs: dict mapping variable names to source step IDs (default: {})
          Used for context passing between steps. When a source step completes,
          its result is injected into downstream steps' resolved_inputs.

    Example step:
      {"id": "build", "type": "task", "title": "Build the module",
       "depends_on": [], "config": {"skills": ["python"], "priority": "high"}}

    Example step with inputs:
      {"id": "review", "type": "task", "title": "Review analysis",
       "inputs": {"analysis": "analyze"},
       "config": {"prompt": "Review this analysis:\\n{analysis}"}}
    """
    if not name:
        return {"error": "Workflow name is required"}

    workflow_id = uuid.uuid4().hex[:12]
    now = datetime.now(timezone.utc).isoformat()

    # Validate steps
    validated_steps: list[dict] = []
    if steps:
        step_ids: set[str] = set()
        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)

            stype = step.get("type", "task")
            if stype not in VALID_STEP_TYPES:
                return {"error": f"Invalid step type '{stype}' for step '{sid}'"}

            validated_steps.append({
                "id": sid,
                "type": stype,
                "title": step.get("title", sid),
                "depends_on": step.get("depends_on", []),
                "config": step.get("config", {}),
                "inputs": step.get("inputs", {}),
            })

        # Validate dependencies exist
        for step in validated_steps:
            for dep in step["depends_on"]:
                if dep not in step_ids:
                    return {"error": f"Step '{step['id']}' depends on unknown step '{dep}'"}

        # Validate input source steps exist
        for step in validated_steps:
            for var_name, source_step_id in step.get("inputs", {}).items():
                if source_step_id not in step_ids:
                    return {"error": f"Step '{step['id']}' input '{var_name}' references unknown step '{source_step_id}'"}

    workflow = {
        "workflow_id": workflow_id,
        "name": name,
        "description": description,
        "steps": validated_steps,
        "created_at": now,
        "updated_at": now,
        "run_count": 0,
    }

    _write_json(DEFINITIONS_DIR / f"{workflow_id}.json", workflow)
    return workflow


def get_workflow(workflow_id: str) -> Optional[dict]:
    """Get a workflow definition."""
    try:
        return json.loads((DEFINITIONS_DIR / f"{workflow_id}.json").read_text("utf-8"))
    except (FileNotFoundError, json.JSONDecodeError, OSError):
        return None


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


def delete_workflow(workflow_id: str) -> bool:
    """Delete a workflow definition."""
    path = DEFINITIONS_DIR / f"{workflow_id}.json"
    if path.exists():
        path.unlink(missing_ok=True)
        return True
    return False


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


def start_workflow(workflow_id: str, original_input: str = "") -> dict:
    """Start a new run of a workflow.

    Creates a run record with all steps initialised to 'pending'.
    Steps with no dependencies are immediately marked 'ready'.
    """
    workflow = get_workflow(workflow_id)
    if not workflow:
        return {"error": f"Workflow {workflow_id} not found"}

    run_id = uuid.uuid4().hex[:12]
    now = datetime.now(timezone.utc).isoformat()

    step_statuses: dict[str, dict] = {}
    for step in workflow["steps"]:
        has_deps = bool(step.get("depends_on"))
        step_statuses[step["id"]] = {
            "status": "pending" if has_deps else "ready",
            "started_at": None,
            "completed_at": None,
            "result": None,
            "error": None,
            "resolved_inputs": {},
        }

    run = {
        "run_id": run_id,
        "workflow_id": workflow_id,
        "workflow_name": workflow["name"],
        "status": "running",
        "original_input": original_input,
        "step_statuses": step_statuses,
        "started_at": now,
        "completed_at": None,
        "steps_total": len(workflow["steps"]),
        "steps_completed": 0,
    }

    _write_json(RUNS_DIR / f"{run_id}.json", run)

    # Update run count
    workflow["run_count"] = workflow.get("run_count", 0) + 1
    workflow["updated_at"] = now
    _write_json(DEFINITIONS_DIR / f"{workflow_id}.json", workflow)

    return run


def complete_step(run_id: str, step_id: str, result: str = "") -> dict:
    """Mark a workflow step as completed and advance the workflow.

    Automatically marks downstream steps as 'ready' when all their
    dependencies are complete. Routes step results to downstream steps
    that declared the completed step as an input source.
    """
    run = get_run(run_id)
    if not run:
        return {"error": f"Run {run_id} not found"}

    step_statuses = run.get("step_statuses", {})
    if step_id not in step_statuses:
        return {"error": f"Step {step_id} not found in run {run_id}"}

    step_statuses[step_id]["status"] = "completed"
    step_statuses[step_id]["completed_at"] = datetime.now(timezone.utc).isoformat()
    step_statuses[step_id]["result"] = result

    run["steps_completed"] = sum(
        1 for s in step_statuses.values() if s["status"] == "completed"
    )

    # Route result to downstream steps that declared this step as an input source
    workflow = get_workflow(run["workflow_id"])
    if workflow:
        for step in workflow["steps"]:
            sid = step["id"]
            step_inputs = step.get("inputs", {})
            for var_name, source_step_id in step_inputs.items():
                if source_step_id == step_id:
                    if "resolved_inputs" not in step_statuses[sid]:
                        step_statuses[sid]["resolved_inputs"] = {}
                    step_statuses[sid]["resolved_inputs"][var_name] = result

    # Advance: check if any pending steps now have all deps completed
    if workflow:
        for step in workflow["steps"]:
            sid = step["id"]
            if step_statuses.get(sid, {}).get("status") != "pending":
                continue
            deps = step.get("depends_on", [])
            if all(step_statuses.get(d, {}).get("status") == "completed" for d in deps):
                step_statuses[sid]["status"] = "ready"

    # Check if all steps complete
    all_done = all(
        s["status"] in ("completed", "skipped") for s in step_statuses.values()
    )
    if all_done:
        run["status"] = "completed"
        run["completed_at"] = datetime.now(timezone.utc).isoformat()

    run["step_statuses"] = step_statuses
    _write_json(RUNS_DIR / f"{run_id}.json", run)

    return {
        "run_id": run_id,
        "step_id": step_id,
        "workflow_status": run["status"],
        "steps_completed": run["steps_completed"],
        "steps_total": run["steps_total"],
        "newly_ready": [
            sid for sid, s in step_statuses.items() if s["status"] == "ready"
        ],
    }


def resolve_step_prompt(run_id: str, step_id: str) -> dict:
    """Resolve a step's prompt template by substituting variable placeholders.

    Reads the step's ``config.prompt`` template and replaces ``{var_name}``
    placeholders with corresponding values from ``resolved_inputs``.
    The special variable ``{original_input}`` is replaced with the run's
    ``original_input`` value.

    Returns a dict with ``prompt`` (the resolved string) on success,
    or ``error`` on failure.
    """
    run = get_run(run_id)
    if not run:
        return {"error": f"Run {run_id} not found"}

    step_statuses = run.get("step_statuses", {})
    if step_id not in step_statuses:
        return {"error": f"Step {step_id} not found in run {run_id}"}

    # Look up the step definition to get config.prompt
    workflow = get_workflow(run["workflow_id"])
    if not workflow:
        return {"error": f"Workflow {run['workflow_id']} not found"}

    step_def = None
    for step in workflow["steps"]:
        if step["id"] == step_id:
            step_def = step
            break
    if not step_def:
        return {"error": f"Step {step_id} not found in workflow definition"}

    prompt_template = step_def.get("config", {}).get("prompt", "")
    if not prompt_template:
        return {"error": f"Step {step_id} has no prompt template in config"}

    # Build substitution map: resolved_inputs + original_input
    resolved = dict(step_statuses[step_id].get("resolved_inputs", {}))
    resolved["original_input"] = run.get("original_input", "")

    # Also support {step_N_output} by scanning all step results
    for sid, sdata in step_statuses.items():
        if sdata.get("result") is not None:
            resolved[f"{sid}_output"] = sdata["result"]

    # Substitute {var_name} placeholders
    resolved_prompt = prompt_template
    for var_name, value in resolved.items():
        placeholder = "{" + var_name + "}"
        resolved_prompt = resolved_prompt.replace(placeholder, str(value))

    return {"prompt": resolved_prompt, "step_id": step_id, "run_id": run_id}


def fail_step(run_id: str, step_id: str, error: str) -> dict:
    """Mark a step as failed. Fails the entire workflow."""
    run = get_run(run_id)
    if not run:
        return {"error": f"Run {run_id} not found"}

    step_statuses = run.get("step_statuses", {})
    if step_id not in step_statuses:
        return {"error": f"Step {step_id} not found"}

    step_statuses[step_id]["status"] = "failed"
    step_statuses[step_id]["error"] = error
    step_statuses[step_id]["completed_at"] = datetime.now(timezone.utc).isoformat()

    run["status"] = "failed"
    run["step_statuses"] = step_statuses
    _write_json(RUNS_DIR / f"{run_id}.json", run)

    return {"run_id": run_id, "step_id": step_id, "status": "failed"}


def skip_step(run_id: str, step_id: str, reason: str = "") -> dict:
    """Skip a step (e.g., conditional branch not taken)."""
    run = get_run(run_id)
    if not run:
        return {"error": f"Run {run_id} not found"}

    step_statuses = run.get("step_statuses", {})
    if step_id not in step_statuses:
        return {"error": f"Step {step_id} not found"}

    step_statuses[step_id]["status"] = "skipped"
    step_statuses[step_id]["result"] = reason
    run["step_statuses"] = step_statuses

    # Advance downstream
    workflow = get_workflow(run["workflow_id"])
    if workflow:
        for step in workflow["steps"]:
            sid = step["id"]
            if step_statuses.get(sid, {}).get("status") != "pending":
                continue
            deps = step.get("depends_on", [])
            if all(step_statuses.get(d, {}).get("status") in ("completed", "skipped") for d in deps):
                step_statuses[sid]["status"] = "ready"

    all_done = all(s["status"] in ("completed", "skipped") for s in step_statuses.values())
    if all_done:
        run["status"] = "completed"
        run["completed_at"] = datetime.now(timezone.utc).isoformat()

    _write_json(RUNS_DIR / f"{run_id}.json", run)
    return {"run_id": run_id, "step_id": step_id, "status": "skipped"}


def cancel_run(run_id: str) -> dict:
    """Cancel a running workflow."""
    run = get_run(run_id)
    if not run:
        return {"error": f"Run {run_id} not found"}
    run["status"] = "cancelled"
    run["completed_at"] = datetime.now(timezone.utc).isoformat()
    _write_json(RUNS_DIR / f"{run_id}.json", run)
    return {"run_id": run_id, "status": "cancelled"}


# ── Queries ────────────────────────────────────────────────────────────────


def get_run(run_id: str) -> Optional[dict]:
    """Get a workflow run."""
    try:
        return json.loads((RUNS_DIR / f"{run_id}.json").read_text("utf-8"))
    except (FileNotFoundError, json.JSONDecodeError, OSError):
        return None


def list_runs(
    workflow_id: Optional[str] = None,
    status: Optional[str] = None,
    limit: int = 50,
) -> list[dict]:
    """List workflow runs, optionally filtered by workflow_id and/or status."""
    runs = []
    for f in sorted(RUNS_DIR.glob("*.json"), reverse=True):
        try:
            r = json.loads(f.read_text("utf-8"))
        except (json.JSONDecodeError, OSError):
            continue
        if workflow_id and r.get("workflow_id") != workflow_id:
            continue
        if status and r.get("status") != status:
            continue
        runs.append(r)
        if len(runs) >= limit:
            break
    return runs


def get_ready_steps(run_id: str) -> list[str]:
    """Get step IDs that are ready to execute in a run."""
    run = get_run(run_id)
    if not run:
        return []
    return [
        sid for sid, s in run.get("step_statuses", {}).items()
        if s["status"] == "ready"
    ]
