"""
Workflow Executor -- Automated execution of multi-step workflow DAGs.

Monitors running workflow runs for "ready" steps, resolves their prompt
templates with variable substitution, dispatches them (direct LLM call
or spawned agent task), and routes results back to the workflow engine
to unlock downstream steps.

THE LOOP:
  1. Scan all running workflow runs for steps with status "ready"
  2. For each ready step:
     a. Resolve the prompt template (variable substitution from upstream results)
     b. Dispatch: direct LLM call (simple steps) or create a broker task (complex)
     c. On completion, call workflow.complete_step() to advance the DAG
     d. On failure, call workflow.fail_step() to halt the run
  3. Sleep, repeat

Usage:
  # Start as background loop (called from dashboard lifespan or CLI)
  python -m brynq.workflow_executor

  # Programmatic
  from brynq.workflow_executor import WorkflowExecutor
  executor = WorkflowExecutor(backend="ollama", model="llama3")
  executor.run()  # blocks
"""

import argparse
import json
import logging
import os
import time
import uuid
from datetime import datetime, timezone
from pathlib import Path
from typing import Callable, Optional

from .server import BROKER_DIR, _write_json_sync
from .workflow import (
    RUNS_DIR,
    complete_step,
    fail_step,
    get_ready_steps,
    get_run,
    get_workflow,
    list_runs,
    resolve_step_prompt,
)

# ── Logging ───────────────────────────────────────────────────────────────

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
    datefmt="%H:%M:%S",
)
log = logging.getLogger("brynq.workflow_executor")

# ── Task directory (same as tasks.py / agent_loop.py) ─────────────────────

TASKS_DIR = BROKER_DIR / "tasks"
TASKS_DIR.mkdir(parents=True, exist_ok=True)


class WorkflowExecutor:
    """Monitors workflow runs and automatically executes ready steps.

    Supports two dispatch modes:
      - **direct**: Send the resolved prompt to an LLM via llm_router and
        feed the response back as the step result. Good for simple chains.
      - **task**: Create a broker task (with ``workflow_run_id`` and ``step_id``
        metadata) and let an agent pick it up. The agent_loop will route the
        result back to the workflow engine on completion.

    Args:
        backend:      LLM backend for direct dispatch (default: "auto").
        model:        Model name for direct dispatch (default: "auto").
        poll_interval: Seconds between scan cycles (default: 5.0).
        dispatch_mode: "direct" (LLM call) or "task" (create broker task).
        llm_fn:       Optional override for the LLM call function.
                      Signature: ``(messages: list[dict]) -> str``.
                      If provided, ``backend`` and ``model`` are ignored.
    """

    def __init__(
        self,
        backend: str = "auto",
        model: str = "auto",
        poll_interval: float = 5.0,
        dispatch_mode: str = "direct",
        llm_fn: Optional[Callable[[list[dict]], str]] = None,
    ) -> None:
        self.backend = backend
        self.model = model
        self.poll_interval = poll_interval
        self.dispatch_mode = dispatch_mode
        self._llm_fn = llm_fn
        self._running = False
        self._iteration = 0
        # Track steps we have already dispatched (to avoid double-dispatch)
        self._dispatched: set[tuple[str, str]] = set()  # (run_id, step_id)

    # ── Public API ────────────────────────────────────────────────────────

    def run(self, max_iterations: Optional[int] = None) -> None:
        """Main executor loop -- scan, dispatch, repeat.

        Args:
            max_iterations: Stop after N cycles. None = run forever.
        """
        self._running = True
        log.info(
            "WorkflowExecutor started (mode=%s, backend=%s, model=%s, poll=%.1fs)",
            self.dispatch_mode, self.backend, self.model, self.poll_interval,
        )

        try:
            while self._running:
                if max_iterations is not None and self._iteration >= max_iterations:
                    break
                self._tick()
                self._iteration += 1
                time.sleep(self.poll_interval)
        except KeyboardInterrupt:
            log.info("WorkflowExecutor stopped (Ctrl+C).")
        finally:
            self._running = False
            log.info("WorkflowExecutor exited after %d iterations.", self._iteration)

    def stop(self) -> None:
        """Signal the executor to stop after the current tick."""
        self._running = False

    def tick_once(self) -> int:
        """Run a single scan-and-dispatch cycle.

        Returns the number of steps dispatched in this tick.
        Useful for testing without running the full loop.
        """
        return self._tick()

    # ── Core Loop ─────────────────────────────────────────────────────────

    def _tick(self) -> int:
        """Scan all running workflow runs and dispatch ready steps.

        Returns the count of steps dispatched.
        """
        dispatched_count = 0

        # Find all running workflow runs
        running_runs = list_runs(status="running")
        if not running_runs:
            return 0

        for run in running_runs:
            run_id = run.get("run_id", "")
            if not run_id:
                continue

            # Re-read to get latest state (another tick may have advanced it)
            fresh_run = get_run(run_id)
            if not fresh_run or fresh_run.get("status") != "running":
                continue

            ready_step_ids = get_ready_steps(run_id)
            for step_id in ready_step_ids:
                dispatch_key = (run_id, step_id)
                if dispatch_key in self._dispatched:
                    continue  # Already dispatched, waiting for result

                try:
                    self._dispatch_step(run_id, step_id, fresh_run)
                    dispatched_count += 1
                except Exception as exc:
                    log.error(
                        "Error dispatching step '%s' in run %s: %s",
                        step_id, run_id, exc,
                    )
                    # Fail the step so the workflow does not hang
                    try:
                        fail_step(run_id, step_id, f"Dispatch error: {exc}")
                    except Exception:
                        pass

        # Clean up dispatched tracking for completed/failed/cancelled runs
        self._cleanup_dispatched()

        if dispatched_count > 0:
            log.info("Dispatched %d step(s) this tick.", dispatched_count)

        return dispatched_count

    # ── Dispatch ──────────────────────────────────────────────────────────

    def _dispatch_step(self, run_id: str, step_id: str, run: dict) -> None:
        """Dispatch a single ready step."""
        if self.dispatch_mode == "task":
            self._dispatch_as_task(run_id, step_id, run)
        else:
            self._dispatch_direct(run_id, step_id, run)

    def _dispatch_direct(self, run_id: str, step_id: str, run: dict) -> None:
        """Execute a step by calling the LLM directly and feeding result back."""
        # Mark step as running in the run file
        self._mark_step_running(run_id, step_id)
        self._dispatched.add((run_id, step_id))

        # Resolve the prompt
        prompt = self._resolve_prompt(run_id, step_id, run)

        # Call LLM
        try:
            result = self._call_llm(prompt)
        except Exception as exc:
            fail_step(run_id, step_id, f"LLM error: {exc}")
            self._dispatched.discard((run_id, step_id))
            return

        result_text = result.strip() if result else "(no output)"

        # Feed result back to workflow engine
        advance = complete_step(run_id, step_id, result_text[:2000])
        self._dispatched.discard((run_id, step_id))

        if "error" in advance:
            log.warning(
                "Workflow advance error for %s/%s: %s",
                run_id, step_id, advance["error"],
            )
        else:
            log.info(
                "Step '%s' completed (%d/%d). Newly ready: %s",
                step_id,
                advance.get("steps_completed", 0),
                advance.get("steps_total", 0),
                advance.get("newly_ready", []),
            )

    def _dispatch_as_task(self, run_id: str, step_id: str, run: dict) -> None:
        """Create a broker task for the step so an agent can pick it up.

        The task includes ``workflow_run_id`` and ``step_id`` metadata so
        that the agent_loop routes the result back to the workflow engine.
        """
        self._mark_step_running(run_id, step_id)
        self._dispatched.add((run_id, step_id))

        prompt = self._resolve_prompt(run_id, step_id, run)

        # Look up step definition for title
        workflow = get_workflow(run.get("workflow_id", ""))
        step_title = step_id
        step_config = {}
        if workflow:
            for step_def in workflow.get("steps", []):
                if step_def["id"] == step_id:
                    step_title = step_def.get("title", step_id)
                    step_config = step_def.get("config", {})
                    break

        task_id = uuid.uuid4().hex[:12]
        channel = step_config.get("channel", "tasks")
        priority = step_config.get("priority", "normal")

        task = {
            "task_id": task_id,
            "channel": channel,
            "status": "pending",
            "requester_session": "workflow-executor",
            "requester_name": "workflow-executor@brynq",
            "assignee_session": None,
            "assignee_name": None,
            "title": f"[WF] {step_title}",
            "description": prompt,
            "created": datetime.now(timezone.utc).isoformat(),
            "accepted_at": None,
            "completed_at": None,
            "result": None,
            "priority": priority,
            "workflow_run_id": run_id,
            "step_id": step_id,
        }

        # Write the task file
        path = TASKS_DIR / f"{task_id}.json"
        tmp = path.with_suffix(".tmp")
        tmp.write_text(json.dumps(task, indent=2, default=str), encoding="utf-8")
        os.replace(str(tmp), str(path))

        log.info(
            "Created task %s for workflow step '%s' (run %s) on #%s",
            task_id, step_id, run_id, channel,
        )

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

    def _resolve_prompt(self, run_id: str, step_id: str, run: dict) -> str:
        """Resolve the prompt for a step, falling back to step title + description."""
        resolved = resolve_step_prompt(run_id, step_id)
        if "prompt" in resolved:
            return resolved["prompt"]

        # Fallback: use step title/description from the workflow definition
        workflow = get_workflow(run.get("workflow_id", ""))
        if workflow:
            for step_def in workflow.get("steps", []):
                if step_def["id"] == step_id:
                    title = step_def.get("title", step_id)
                    config = step_def.get("config", {})
                    desc = config.get("description", "")
                    # Also inject any resolved inputs as context
                    step_statuses = run.get("step_statuses", {})
                    inputs = step_statuses.get(step_id, {}).get("resolved_inputs", {})
                    parts = [f"Task: {title}"]
                    if desc:
                        parts.append(f"Description: {desc}")
                    if inputs:
                        parts.append("Context from previous steps:")
                        for var, val in inputs.items():
                            parts.append(f"  {var}: {val}")
                    return "\n".join(parts)

        return f"Execute workflow step: {step_id}"

    def _call_llm(self, prompt: str) -> str:
        """Call the LLM with a prompt. Uses llm_fn override if provided."""
        messages = [{"role": "user", "content": prompt}]

        if self._llm_fn is not None:
            return self._llm_fn(messages)

        from . import llm_router
        return llm_router.chat(self.backend, self.model, messages)

    def _mark_step_running(self, run_id: str, step_id: str) -> None:
        """Update the step status to 'running' in the run file."""
        run = get_run(run_id)
        if not run:
            return
        step_statuses = run.get("step_statuses", {})
        if step_id in step_statuses:
            step_statuses[step_id]["status"] = "running"
            step_statuses[step_id]["started_at"] = datetime.now(timezone.utc).isoformat()
            run["step_statuses"] = step_statuses
            _write_json_sync(RUNS_DIR / f"{run_id}.json", run)

    def _cleanup_dispatched(self) -> None:
        """Remove dispatch tracking entries for runs that are no longer running."""
        if not self._dispatched:
            return
        stale = set()
        for run_id, step_id in self._dispatched:
            run = get_run(run_id)
            if not run or run.get("status") != "running":
                stale.add((run_id, step_id))
                continue
            # Also clean up if step is no longer in running/ready state
            step_status = run.get("step_statuses", {}).get(step_id, {}).get("status")
            if step_status in ("completed", "failed", "skipped", None):
                stale.add((run_id, step_id))
        self._dispatched -= stale


# ── CLI Entry Point ───────────────────────────────────────────────────────


def main() -> None:
    """Parse arguments and start the workflow executor."""
    parser = argparse.ArgumentParser(
        prog="brynq.workflow_executor",
        description=(
            "Run the Brynq Workflow Executor. Monitors running workflow "
            "runs and automatically dispatches ready steps to LLMs or agents."
        ),
    )
    parser.add_argument(
        "--backend",
        default="auto",
        help="LLM backend for direct dispatch (default: auto)",
    )
    parser.add_argument(
        "--model",
        default="auto",
        help="Model name for direct dispatch (default: auto)",
    )
    parser.add_argument(
        "--poll-interval",
        type=float,
        default=5.0,
        help="Seconds between scan cycles (default: 5.0)",
    )
    parser.add_argument(
        "--mode",
        default="direct",
        choices=["direct", "task"],
        help="Dispatch mode: 'direct' (LLM call) or 'task' (broker task) (default: direct)",
    )
    parser.add_argument(
        "--max-iterations",
        type=int,
        default=None,
        help="Stop after N iterations (default: run forever)",
    )

    args = parser.parse_args()

    executor = WorkflowExecutor(
        backend=args.backend,
        model=args.model,
        poll_interval=args.poll_interval,
        dispatch_mode=args.mode,
    )
    executor.run(max_iterations=args.max_iterations)


if __name__ == "__main__":
    main()
