"""Knowledge ingestion helper."""

from __future__ import annotations

import time
from typing import Any, Dict, Optional

from ..client import _HTTPClient


FINAL_STATES = {"succeeded", "failed"}


class KnowledgeAPI:
    def __init__(self, http: _HTTPClient) -> None:
        self._http = http

    def ingest(
        self,
        companion_id: str,
        *,
        payload_type: str = "text",
        content: Optional[str] = None,
        key: Optional[str] = None,
    ) -> dict[str, Any]:
        body: Dict[str, Any] = {"type": payload_type}
        if content is not None:
            body["content"] = content
        if key is not None:
            body["key"] = key
        response = self._http.request(
            "POST",
            f"/v1/companions/{companion_id}/knowledge",
            json=body,
            expected=(202,),
        )
        if not isinstance(response, dict):  # pragma: no cover - defensive
            raise TypeError("Expected dict payload for knowledge ingestion response")
        return response

    def get_job(self, job_id: str) -> dict[str, Any]:
        job = self._http.request("GET", f"/v1/knowledge-jobs/{job_id}")
        if not isinstance(job, dict):  # pragma: no cover - defensive
            raise TypeError("Expected dict payload for knowledge job")
        return job

    def wait(self, job_id: str, *, timeout: float = 30.0, interval: float = 0.5) -> dict[str, Any]:
        deadline = time.monotonic() + timeout
        while True:
            job = self.get_job(job_id)
            if job.get("status") in FINAL_STATES:
                return job
            if time.monotonic() >= deadline:
                return job
            time.sleep(interval)
