#!/usr/bin/env python3
"""Keep talking to a running Ask session (mid-turn steer).

You start an investigation, then send a follow-up on the same session_id
(deploy window, better log line, Slack reply). When the running turn takes it,
the start reply sets action=\"steer\" and you wait on that one trace_id instead
of opening a parallel dig.

Requires STACKGEN_URL, STACKGEN_TOKEN, STACKGEN_WORKSPACE_ID.

Usage:
  export STACKGEN_URL=https://app.stackgen.com
  export STACKGEN_TOKEN=stackgen_…
  export STACKGEN_WORKSPACE_ID=<workspace-uuid>
  python examples/ask/ask-steer-followup.py --message "Investigate checkout latency"
"""

from __future__ import annotations

import argparse
import os
import sys
import uuid

from stackgen import StackgenClient, StackgenConfig, is_steer_response


def env(name: str, *fallbacks: str) -> str:
    for key in (name, *fallbacks):
        value = (os.environ.get(key) or "").strip()
        if value:
            return value
    raise SystemExit(f"missing env {name}")


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--message",
        default="Investigate a synthetic checkout latency spike",
        help="Initial Ask message",
    )
    parser.add_argument(
        "--follow-up",
        default="Also check recent deploy markers before you finish",
        help="Mid-session follow-up text",
    )
    parser.add_argument(
        "--workflow",
        default=os.environ.get("STACKGEN_WORKFLOW", ""),
        help="Optional workflow entity_ref (CI fixture)",
    )
    parser.add_argument("--timeout-seconds", type=int, default=600)
    args = parser.parse_args()

    client = StackgenClient(
        StackgenConfig(
            base_url=env("STACKGEN_URL"),
            api_token=env("STACKGEN_TOKEN"),
            workspace_id=env("STACKGEN_WORKSPACE_ID", "STACKGEN_PROJECT"),
            timeout_seconds=args.timeout_seconds,
        )
    )
    session_id = f"sdk-steer-{uuid.uuid4().hex[:12]}"
    entity_refs = (
        [{"name": args.workflow, "type": "workflow"}] if args.workflow.strip() else None
    )

    def wait_settled(trace_id: str) -> bool:
        try:
            done = client.aiden.wait_for_subscribe_complete(
                trace_id, timeout_ms=args.timeout_seconds * 1000
            )
            if done.ok:
                print(f"subscribe complete ok=True trace_id={trace_id}")
                return True
            print(
                f"subscribe incomplete ({done}); waiting on execution status",
                file=sys.stderr,
            )
        except Exception as exc:  # noqa: BLE001
            message = str(exc)
            if not any(
                token in message.lower()
                for token in ("timed out", "subscribeexecution", "aborted", "eof")
            ):
                raise
            print(
                f"subscribe timed out ({exc}); waiting on execution status",
                file=sys.stderr,
            )
        trace = client.aiden.wait_for_execution(trace_id)
        execution = trace.get("execution") if isinstance(trace, dict) else {}
        if not isinstance(execution, dict):
            execution = {}
        status = str(execution.get("status") or "").lower()
        ok = status in {"completed", "success", "succeeded"}
        print(f"subscribe complete ok={ok} trace_id={trace_id} execution={status}")
        return ok

    try:
        start = client.aiden.start_ask(
            args.message, session_id=session_id, entity_refs=entity_refs
        )
        trace_id = str(start.get("trace_id") or "")
        if not trace_id:
            raise SystemExit(f"start_ask missing trace_id: {start}")
        print(f"started session_id={session_id} trace_id={trace_id}")

        follow = client.aiden.start_ask(
            args.follow_up, session_id=session_id, entity_refs=entity_refs
        )
        follow_trace = str(follow.get("trace_id") or "")
        steered = is_steer_response(follow)
        print(
            f"follow-up action={follow.get('action')!r} steered={steered} "
            f"trace_id={follow_trace}"
        )

        wait_trace = follow_trace or trace_id
        ok = wait_settled(wait_trace)
        if steered and follow_trace and follow_trace != trace_id:
            print(
                "warning: steer response used a different trace_id than the original start",
                file=sys.stderr,
            )
        return 0 if ok else 1
    finally:
        try:
            client.aiden.terminate_session(session_id, reason="sdk example finished")
        except Exception as exc:  # noqa: BLE001
            print(f"terminate session {session_id}: {exc}", file=sys.stderr)


if __name__ == "__main__":
    raise SystemExit(main())