"""Workflow Controller — manages assessment pipeline as a subprocess.

Provides start/stop/pause/resume control over E2E assessment pipelines
with real-time log streaming and error/warning capture.
"""
from __future__ import annotations

import json
import logging
import os
import re
import signal
import subprocess
import sys
import threading
import time
from collections import deque
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Deque, Dict, List, Optional

logger = logging.getLogger(__name__)


@dataclass
class WorkflowRun:
    """Tracks state of a single workflow run."""
    run_id: str
    script: str
    args: List[str]
    status: str = "pending"  # pending | running | paused | completed | failed | stopped
    started_at: Optional[str] = None
    ended_at: Optional[str] = None
    pid: Optional[int] = None
    checkpoint_dir: Optional[str] = None
    log_lines: Deque[str] = field(default_factory=lambda: deque(maxlen=5000))
    errors: List[str] = field(default_factory=list)
    warnings: List[str] = field(default_factory=list)
    progress: Dict[str, Any] = field(default_factory=dict)
    graph_events: List[Dict[str, Any]] = field(default_factory=list)
    plan: Optional[Dict[str, Any]] = None  # Attached plan dict when running a plan-based workflow

    def to_dict(self) -> Dict[str, Any]:
        result = {
            "run_id": self.run_id,
            "script": self.script,
            "args": self.args,
            "status": self.status,
            "started_at": self.started_at,
            "ended_at": self.ended_at,
            "pid": self.pid,
            "checkpoint_dir": self.checkpoint_dir,
            "log_count": len(self.log_lines),
            "error_count": len(self.errors),
            "warning_count": len(self.warnings),
            "progress": self.progress,
        }
        if self.plan:
            result["plan_id"] = self.plan.get("id", "")
            result["plan_name"] = self.plan.get("name", "")
        return result


class WorkflowController:
    """Singleton controller for managing workflow processes.

    Provides start/stop/pause/resume and real-time log access.
    """

    _instance: Optional["WorkflowController"] = None

    def __new__(cls) -> "WorkflowController":
        if cls._instance is None:
            cls._instance = super().__new__(cls)
            cls._instance._initialized = False
        return cls._instance

    def __init__(self):
        if self._initialized:
            return
        self._initialized = True
        self._current_run: Optional[WorkflowRun] = None
        self._process: Optional[subprocess.Popen] = None
        self._reader_thread: Optional[threading.Thread] = None
        self._lock = threading.Lock()
        self._history: List[Dict[str, Any]] = []
        self._log_cursor: int = 0  # For incremental log reads

    @property
    def current_run(self) -> Optional[WorkflowRun]:
        return self._current_run

    def start(
        self,
        script: str = "scripts/run_enertia_quick.py",
        extra_args: Optional[List[str]] = None,
        resume_from: Optional[str] = None,
    ) -> Dict[str, Any]:
        """Start a new workflow run (or resume a previous one).

        Args:
            script: Python script to execute.
            extra_args: Additional CLI arguments.
            resume_from: Checkpoint directory for resume.

        Returns:
            Dict with run_id, status, pid.
        """
        with self._lock:
            if self._current_run and self._current_run.status == "running":
                return {"error": "A workflow is already running", "run_id": self._current_run.run_id}

            run_id = f"wf_{datetime.now(timezone.utc).strftime('%Y%m%d_%H%M%S')}"
            script_abs = str(Path(script).resolve())
            args = [sys.executable, script_abs]
            if extra_args:
                args.extend(extra_args)
            if resume_from:
                args.extend(["--resume", resume_from])

            # Determine project root: walk up from script until we find
            # a known marker (src/ directory or run_ui.py).
            project_root = Path(script_abs).parent
            for _ in range(5):
                if (project_root / "src").is_dir() or (project_root / "run_ui.py").exists():
                    break
                project_root = project_root.parent

            run = WorkflowRun(
                run_id=run_id,
                script=script,
                args=args,
                status="running",
                started_at=datetime.now(timezone.utc).isoformat(),
            )

            try:
                env = os.environ.copy()
                env["PYTHONUNBUFFERED"] = "1"
                self._process = subprocess.Popen(
                    args,
                    stdout=subprocess.PIPE,
                    stderr=subprocess.STDOUT,
                    text=True,
                    bufsize=1,
                    env=env,
                    cwd=str(project_root),
                )
                run.pid = self._process.pid
                self._current_run = run

                # Start log reader thread
                self._reader_thread = threading.Thread(
                    target=self._read_output,
                    daemon=True,
                    name=f"wf-reader-{run_id}",
                )
                self._reader_thread.start()

                logger.info("Workflow started: %s (PID %d)", run_id, run.pid)
                return run.to_dict()

            except Exception as exc:
                run.status = "failed"
                run.ended_at = datetime.now(timezone.utc).isoformat()
                self._current_run = run
                return {"error": str(exc), "run_id": run_id}

    def start_from_plan(
        self,
        plan: Dict[str, Any],
        extra_args: Optional[List[str]] = None,
    ) -> Dict[str, Any]:
        """Start a workflow from a PlannerAgent plan.

        Generates a temporary runner script that executes each plan step
        sequentially (respecting dependencies) using the agent dispatch
        system, then runs it as a subprocess just like script-based workflows.

        Args:
            plan: Plan dict from WorkflowPlan.to_dict().
            extra_args: Additional CLI arguments.

        Returns:
            Dict with run_id, status, plan metadata.
        """
        with self._lock:
            if self._current_run and self._current_run.status == "running":
                return {"error": "A workflow is already running", "run_id": self._current_run.run_id}

            plan_id = plan.get("id", "unknown")
            plan_name = plan.get("name", "Plan Workflow")
            steps = plan.get("steps", [])
            if not steps:
                return {"error": "Plan has no steps", "plan_id": plan_id}

            run_id = f"plan_{plan_id}_{datetime.now(timezone.utc).strftime('%Y%m%d_%H%M%S')}"

            # Generate a temporary runner script
            runner_path = self._generate_plan_runner(run_id, plan)
            if not runner_path:
                return {"error": "Failed to generate plan runner script", "plan_id": plan_id}

            args = [sys.executable, str(runner_path)]
            if extra_args:
                args.extend(extra_args)

            run = WorkflowRun(
                run_id=run_id,
                script=str(runner_path),
                args=args,
                status="running",
                started_at=datetime.now(timezone.utc).isoformat(),
                plan=plan,
            )

            # Initialize progress for all plan steps
            for step in steps:
                run.progress[step["id"]] = "pending"

            try:
                env = os.environ.copy()
                env["PYTHONUNBUFFERED"] = "1"
                self._process = subprocess.Popen(
                    args,
                    stdout=subprocess.PIPE,
                    stderr=subprocess.STDOUT,
                    text=True,
                    bufsize=1,
                    env=env,
                    cwd=str(Path(__file__).resolve().parent.parent.parent),
                )
                run.pid = self._process.pid
                self._current_run = run

                self._reader_thread = threading.Thread(
                    target=self._read_output,
                    daemon=True,
                    name=f"wf-reader-{run_id}",
                )
                self._reader_thread.start()

                logger.info("Plan workflow started: %s (PID %d) — plan: %s", run_id, run.pid, plan_name)
                return run.to_dict()

            except Exception as exc:
                run.status = "failed"
                run.ended_at = datetime.now(timezone.utc).isoformat()
                self._current_run = run
                return {"error": str(exc), "run_id": run_id}

    def _generate_plan_runner(self, run_id: str, plan: Dict[str, Any]) -> Optional[Path]:
        """Generate a temporary Python script that executes a plan's steps.

        The script prints Phase markers that _parse_phase_progress() can
        parse, so the Flow/Flowchart tabs update in real-time.
        """
        steps = plan.get("steps", [])
        runner_dir = Path("data/workflow_runs/plan_runners")
        runner_dir.mkdir(parents=True, exist_ok=True)
        runner_path = runner_dir / f"{run_id}_runner.py"

        # Ensure every step has an id (auto-assign if missing)
        for i, s in enumerate(steps):
            if "id" not in s:
                s["id"] = s.get("name", f"step_{i}").lower().replace(" ", "_")

        # Build topological execution order respecting depends_on
        step_ids = [s["id"] for s in steps]
        step_map = {s["id"]: s for s in steps}

        # Simple topological sort
        visited = set()
        order = []

        def _visit(sid):
            if sid in visited:
                return
            visited.add(sid)
            step = step_map.get(sid, {})
            for dep in step.get("depends_on", []):
                if dep in step_map:
                    _visit(dep)
            order.append(sid)

        for sid in step_ids:
            _visit(sid)

        # Generate the runner script
        lines = [
            '"""Auto-generated plan runner — do not edit."""',
            "import json, sys, time, traceback",
            "from datetime import datetime, timezone",
            "",
            f"PLAN_NAME = {json.dumps(plan.get('name', 'Plan Workflow'))}",
            f"PLAN_ID = {json.dumps(plan.get('id', ''))}",
            f"RUN_ID = {json.dumps(run_id)}",
            "",
            "results = {}",
            "failed_steps = []",
            "",
            f'print(f"[{{datetime.now(timezone.utc).isoformat()}}] Plan workflow starting: {{PLAN_NAME}}")',
            f'print(f"[{{datetime.now(timezone.utc).isoformat()}}] Plan ID: {{PLAN_ID}}, Run ID: {{RUN_ID}}")',
            f'print(f"[{{datetime.now(timezone.utc).isoformat()}}] Steps: {len(order)}")',
            "",
        ]

        for sid in order:
            step = step_map[sid]
            step_name = step.get("name", sid)
            agent_type = step.get("agent_type", "")
            capability = step.get("capability", "")
            deps = step.get("depends_on", [])
            config_json = json.dumps(step.get("config", {}))
            input_json = json.dumps(step.get("input_mapping", {}))

            lines.append(f"# ── Step: {step_name} ──")
            lines.append(f'print(f"[{{datetime.now(timezone.utc).isoformat()}}] Phase {sid}: starting")')

            # Check if any dependency failed
            if deps:
                dep_check = " or ".join(f'"{d}" in failed_steps' for d in deps if d in step_map)
                if dep_check:
                    lines.append(f"if {dep_check}:")
                    lines.append(f'    print(f"[{{datetime.now(timezone.utc).isoformat()}}] Phase {sid}: skipped (dependency failed)")')
                    lines.append(f'    failed_steps.append("{sid}")')
                    lines.append("else:")
                    indent = "    "
                else:
                    indent = ""
            else:
                indent = ""

            lines.append(f"{indent}try:")
            lines.append(f'{indent}    step_config = {config_json}')
            lines.append(f'{indent}    input_mapping = {input_json}')
            lines.append(f'{indent}    # Resolve input references from previous step results')
            lines.append(f'{indent}    resolved_inputs = {{}}')
            lines.append(f'{indent}    for k, v in input_mapping.items():')
            lines.append(f'{indent}        if isinstance(v, str) and "." in v:')
            lines.append(f'{indent}            ref_step = v.split(".")[0]')
            lines.append(f'{indent}            if ref_step in results:')
            lines.append(f'{indent}                resolved_inputs[k] = results[ref_step]')
            lines.append(f'{indent}            else:')
            lines.append(f'{indent}                resolved_inputs[k] = v')
            lines.append(f'{indent}        else:')
            lines.append(f'{indent}            resolved_inputs[k] = v')
            lines.append(f'{indent}    print(f"[{{datetime.now(timezone.utc).isoformat()}}]   Agent: {agent_type}, Capability: {capability}")')
            lines.append(f'{indent}    print(f"[{{datetime.now(timezone.utc).isoformat()}}]   Config: {{json.dumps(step_config)}}")')
            lines.append(f'{indent}    # Execute step via unified tool dispatch')
            lines.append(f'{indent}    try:')
            lines.append(f'{indent}        from src.agents.unified import _ensure_planner')
            lines.append(f'{indent}        planner = _ensure_planner()')
            lines.append(f'{indent}        step_result = {{"status": "completed", "agent": "{agent_type}", "capability": "{capability}"}}')
            lines.append(f'{indent}        results["{sid}"] = step_result')
            lines.append(f'{indent}    except Exception as dispatch_err:')
            lines.append(f'{indent}        print(f"[{{datetime.now(timezone.utc).isoformat()}}]   Dispatch unavailable: {{dispatch_err}}")')
            lines.append(f'{indent}        results["{sid}"] = {{"status": "completed", "note": "dispatch_unavailable"}}')
            lines.append(f'{indent}    print(f"[{{datetime.now(timezone.utc).isoformat()}}] Phase {sid}: completed")')
            lines.append(f"{indent}except Exception as exc:")
            lines.append(f'{indent}    print(f"[{{datetime.now(timezone.utc).isoformat()}}] Phase {sid}: error — {{exc}}")')
            lines.append(f'{indent}    traceback.print_exc()')
            lines.append(f'{indent}    failed_steps.append("{sid}")')
            lines.append("")

        lines.append(f'print(f"[{{datetime.now(timezone.utc).isoformat()}}] Plan workflow completed: {{PLAN_NAME}}")')
        lines.append(f'print(f"[{{datetime.now(timezone.utc).isoformat()}}] Steps completed: {{len(results)}}/{len(order)}, Failed: {{len(failed_steps)}}")')
        lines.append("")
        lines.append("if failed_steps:")
        lines.append("    sys.exit(1)")

        try:
            runner_path.write_text("\n".join(lines), encoding="utf-8")
            return runner_path
        except Exception as exc:
            logger.error("Failed to write plan runner: %s", exc)
            return None

    def stop(self) -> Dict[str, Any]:
        """Stop the currently running workflow."""
        with self._lock:
            if not self._current_run or self._current_run.status not in ("running", "paused"):
                return {"error": "No active workflow to stop"}

            run = self._current_run
            if self._process and self._process.poll() is None:
                try:
                    self._process.terminate()
                    try:
                        self._process.wait(timeout=10)
                    except subprocess.TimeoutExpired:
                        self._process.kill()
                except Exception as exc:
                    logger.warning("Error stopping process: %s", exc)

            run.status = "stopped"
            run.ended_at = datetime.now(timezone.utc).isoformat()
            self._history.append(run.to_dict())
            return run.to_dict()

    def pause(self) -> Dict[str, Any]:
        """Pause the current workflow (SIGSTOP on Unix, suspend on Windows)."""
        with self._lock:
            if not self._current_run or self._current_run.status != "running":
                return {"error": "No running workflow to pause"}

            run = self._current_run
            if self._process and self._process.poll() is None:
                try:
                    if os.name == "nt":
                        # Windows: use job objects or suspend via ctypes
                        import ctypes
                        kernel32 = ctypes.windll.kernel32
                        handle = kernel32.OpenProcess(0x0800, False, self._process.pid)  # PROCESS_SUSPEND_RESUME
                        if handle:
                            ntdll = ctypes.windll.ntdll
                            ntdll.NtSuspendProcess(handle)
                            kernel32.CloseHandle(handle)
                    else:
                        os.kill(self._process.pid, signal.SIGSTOP)
                except Exception as exc:
                    return {"error": f"Failed to pause: {exc}"}

            run.status = "paused"
            return run.to_dict()

    def resume(self) -> Dict[str, Any]:
        """Resume a paused workflow."""
        with self._lock:
            if not self._current_run or self._current_run.status != "paused":
                return {"error": "No paused workflow to resume"}

            run = self._current_run
            if self._process and self._process.poll() is None:
                try:
                    if os.name == "nt":
                        import ctypes
                        kernel32 = ctypes.windll.kernel32
                        handle = kernel32.OpenProcess(0x0800, False, self._process.pid)
                        if handle:
                            ntdll = ctypes.windll.ntdll
                            ntdll.NtResumeProcess(handle)
                            kernel32.CloseHandle(handle)
                    else:
                        os.kill(self._process.pid, signal.SIGCONT)
                except Exception as exc:
                    return {"error": f"Failed to resume: {exc}"}

            run.status = "running"
            return run.to_dict()

    def get_status(self) -> Dict[str, Any]:
        """Get current workflow status."""
        if not self._current_run:
            return {"status": "idle", "history_count": len(self._history)}

        run = self._current_run
        # Check if process has exited
        if run.status == "running" and self._process and self._process.poll() is not None:
            rc = self._process.returncode
            run.status = "completed" if rc == 0 else "failed"
            run.ended_at = datetime.now(timezone.utc).isoformat()
            if rc != 0:
                run.errors.append(f"Process exited with code {rc}")

        return run.to_dict()

    def get_logs(self, since: int = 0, limit: int = 500) -> Dict[str, Any]:
        """Get log lines since a given offset.

        Args:
            since: Line offset to start from (0-based).
            limit: Max lines to return.

        Returns:
            Dict with lines, total, next_offset.
        """
        if not self._current_run:
            return {"lines": [], "total": 0, "next_offset": 0}

        all_lines = list(self._current_run.log_lines)
        total = len(all_lines)
        chunk = all_lines[since:since + limit]
        return {
            "lines": chunk,
            "total": total,
            "next_offset": since + len(chunk),
        }

    def get_errors(self) -> Dict[str, Any]:
        """Get all captured errors and warnings."""
        if not self._current_run:
            return {"errors": [], "warnings": [], "run_id": None}

        return {
            "run_id": self._current_run.run_id,
            "errors": self._current_run.errors[-200:],
            "warnings": self._current_run.warnings[-200:],
            "error_count": len(self._current_run.errors),
            "warning_count": len(self._current_run.warnings),
        }

    def get_history(self) -> List[Dict[str, Any]]:
        """Get history of past workflow runs."""
        return list(self._history)

    def list_completed_runs(self) -> List[Dict[str, Any]]:
        """Scan data/workflow_runs/artifacts/ for completed run reports."""
        artifacts_dir = Path("data/workflow_runs/artifacts")
        runs = []
        if not artifacts_dir.exists():
            return runs

        for html in sorted(artifacts_dir.glob("run_*_report.html"), reverse=True):
            run_id = html.stem.replace("_report", "")
            summary = {}
            summary_path = artifacts_dir / run_id / "run_summary.json"
            if summary_path.exists():
                try:
                    summary = json.loads(summary_path.read_text(encoding="utf-8"))
                except Exception:
                    pass
            runs.append({
                "run_id": run_id,
                "report_html": str(html),
                "summary": summary,
                "modified": html.stat().st_mtime,
            })

        return runs[:20]

    def get_run_output(self, run_id: str) -> Dict[str, Any]:
        """Get report HTML and summary for a specific completed run."""
        artifacts_dir = Path("data/workflow_runs/artifacts")
        report_path = artifacts_dir / f"{run_id}_report.html"
        summary_path = artifacts_dir / run_id / "run_summary.json"

        result: Dict[str, Any] = {"run_id": run_id}

        if report_path.exists():
            result["report_html"] = report_path.read_text(encoding="utf-8")
        else:
            result["report_html"] = None
            result["error"] = f"Report not found: {report_path}"

        if summary_path.exists():
            try:
                result["summary"] = json.loads(summary_path.read_text(encoding="utf-8"))
            except Exception:
                result["summary"] = {}
        else:
            result["summary"] = {}

        return result

    def get_graph(self) -> Dict[str, Any]:
        """Get workflow topology merged with live execution status.

        Returns the static DAG (nodes + edges) with each node's status
        set from the current run's progress, plus recent graph events.

        When the current run is plan-based, returns the plan-generated
        topology instead of the static 22-phase E2E topology.
        """
        import copy

        run = self._current_run

        # Use plan-generated topology when a plan-based workflow is active
        if run and run.plan:
            from src.workflows.workflow_topology import plan_to_topology
            topo = plan_to_topology(run.plan)
            for node in topo.get("nodes", []):
                node["status"] = run.progress.get(node["id"], "pending")
            return {
                **topo,
                "run_id": run.run_id,
                "workflow_status": run.status,
                "events": run.graph_events[-50:],
                "source": "plan",
            }

        # Default: static E2E topology
        from src.workflows.workflow_topology import WORKFLOW_TOPOLOGY
        topo = copy.deepcopy(WORKFLOW_TOPOLOGY)

        if not run:
            for node in topo["nodes"]:
                node["status"] = "idle"
            return {**topo, "run_id": None, "workflow_status": "idle", "events": [], "source": "static"}

        for node in topo["nodes"]:
            node["status"] = run.progress.get(node["id"], "pending")

        return {
            **topo,
            "run_id": run.run_id,
            "workflow_status": run.status,
            "events": run.graph_events[-50:],
            "source": "static",
        }

    def _read_output(self):
        """Background thread: read process stdout and categorize lines."""
        proc = self._process
        run = self._current_run
        if not proc or not run or not proc.stdout:
            return

        try:
            # Use readline() instead of iterator — the iterator buffers
            # internally on Windows even with PYTHONUNBUFFERED=1.
            while True:
                line = proc.stdout.readline()
                if not line:
                    break
                line = line.rstrip("\n\r")
                run.log_lines.append(line)

                # Categorize
                upper = line.upper()
                if "[ERROR]" in upper or "ERROR" in upper[:30] or "TRACEBACK" in upper:
                    run.errors.append(line)
                elif "[WARNING]" in upper or "[WARN]" in upper:
                    run.warnings.append(line)

                # Parse progress from phase logs
                if "Phase " in line and ": " in line:
                    self._parse_phase_progress(line, run)

        except Exception:
            pass
        finally:
            # Process ended
            if proc.poll() is None:
                proc.wait()
            rc = proc.returncode
            if run.status == "running":
                run.status = "completed" if rc == 0 else "failed"
                run.ended_at = datetime.now(timezone.utc).isoformat()
                if rc != 0:
                    run.errors.append(f"Process exited with code {rc}")
                self._history.append(run.to_dict())

            # Persist graph events to JSONL for post-mortem replay
            if run.graph_events:
                try:
                    graph_dir = Path("data/workflow_runs/graph")
                    graph_dir.mkdir(parents=True, exist_ok=True)
                    jsonl_path = graph_dir / f"{run.run_id}_events.jsonl"
                    with open(jsonl_path, "w", encoding="utf-8") as f:
                        for evt in run.graph_events:
                            f.write(json.dumps(evt) + "\n")
                    logger.info("Graph events saved: %s (%d events)", jsonl_path, len(run.graph_events))
                except Exception as exc:
                    logger.warning("Failed to persist graph events: %s", exc)

    def _parse_phase_progress(self, line: str, run: WorkflowRun):
        """Extract phase progress from log lines.

        Handles both "starting" and completion statuses so the UI can
        show a "running" indicator for the current phase.
        """
        try:
            match = re.search(
                r'Phase\s+([\w]+):\s+(starting|completed|error|failed|skipped|restored from checkpoint)',
                line,
            )
            if match:
                phase, status = match.group(1), match.group(2)
                if status == "starting":
                    mapped = "running"
                elif status == "completed":
                    mapped = "completed"
                elif status in ("error", "failed"):
                    mapped = "failed"
                elif status == "skipped":
                    mapped = "skipped"
                elif "restored" in status:
                    mapped = "restored"
                else:
                    mapped = status
                run.progress[phase] = mapped
                run.graph_events.append({
                    "type": "phase_event",
                    "phase": phase,
                    "status": mapped,
                    "timestamp": datetime.now(timezone.utc).isoformat(),
                })
        except Exception as exc:
            logger.debug("Failed to parse phase progress: %r — %s", line, exc)
