#!/usr/bin/env python3
"""Set a goal and observe automatic rounds through the final model reply."""

from __future__ import annotations

import argparse
import math
import os
import time
from pathlib import Path

from pydantic import BaseModel

from deepseek_harness import CommandExecution, DeepSeekHarness, Notification
from deepseek_harness.api import final_response


class WaitResult(BaseModel):
    """The empty session/wait result."""


def display(notification: Notification) -> dict | None:
    """Print durable assistant output and return the observed root Session event."""
    if notification.method != "session.event":
        return None
    event = notification.payload["event"]
    if event["type"] == "assistant/message":
        text = final_response([event])
        if text:
            print(text, flush=True)
    elif event["type"] == "goal/change":
        print(f"Goal: {event['data']['operation']}", flush=True)
    elif event["type"] == "turn/end" and event["data"]["reason"]["kind"] != "completed":
        raise RuntimeError(f"Goal activity ended: {event['data']['reason']}")
    return event


def main() -> None:
    """Observe goal state before waiting for the final activity interval."""
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--dsh-home", required=True, type=Path)
    parser.add_argument("--workspace", type=Path, default=Path.cwd())
    parser.add_argument("--model", default=os.environ.get("DSH_MODEL", "deepseek-v4-flash"))
    parser.add_argument("--timeout-seconds", type=float, default=300)
    parser.add_argument("objective", help="Objective for the native /goal command")
    args = parser.parse_args()
    if not args.objective.strip() or not math.isfinite(args.timeout_seconds) or args.timeout_seconds <= 0:
        parser.error("objective must be non-blank and timeout-seconds must be finite and positive")
    with DeepSeekHarness(
        dsh_home=str(args.dsh_home.resolve()), cwd=str(args.workspace.resolve()),
        profile="sdk", model=args.model, env={"DSH_PERMISSION_MODE": "workspace-write"},
        request_timeout_seconds=args.timeout_seconds,
    ) as harness:
        session = harness.start_session()
        with harness.client.subscribe_notifications(lambda notice: notice.payload.get("sessionId") == session.id) as notifications:
            deadline = time.monotonic() + args.timeout_seconds

            def remaining_seconds() -> float:
                remaining = deadline - time.monotonic()
                if remaining <= 0:
                    raise TimeoutError("Goal observation timed out")
                return remaining

            def observe(notification: Notification) -> dict | None:
                remaining_seconds()
                event = display(notification)
                remaining_seconds()
                return event

            execution = session.run(f"/goal {args.objective}")
            if not isinstance(execution, CommandExecution):
                raise RuntimeError("The SDK profile did not execute the requested command")
            print(execution.result.text, flush=True)
            if execution.result.kind == "error":
                raise RuntimeError(execution.result.text)
            while True:
                event = observe(notifications.next(timeout_seconds=remaining_seconds()))
                if event is not None and event["type"] == "goal/change" and event["data"]["operation"] in {
                    "complete", "block", "pause", "clear",
                }:
                    break
            # A terminal goal event can precede the model's closing reply.
            harness.client.request(
                "session/wait", {"sessionId": session.id}, response_model=WaitResult,
                timeout_seconds=remaining_seconds(),
                on_notification=observe, notification_subscription=notifications,
            )
            remaining_seconds()


if __name__ == "__main__":
    main()