# 文件路径及文件名: laybot_ai/batch.py
# 简要阐述该文件的作用: 封装 OpenAI Batch 的上传、创建、查询、取消、等待完成、结果下载与结果解析等核心能力。

from __future__ import annotations
import time
from typing import Any

from .client import Client
from .file import File
from .exceptions import LayBotError
from .jsonl import (
    parse_jsonl,
    iter_jsonl,
    iter_jsonl_safe,
    parse_jsonl_file,
    iter_jsonl_file,
    iter_jsonl_file_safe,
    dump_jsonl,
)


class Batch:
    TERMINAL_STATUSES = {"completed", "failed", "cancelled", "expired"}

    def __init__(self, api_key: str | Client, *, vendor="openai", base=None):
        self.cli = api_key if isinstance(api_key, Client) \
            else Client(api_key, vendor=vendor, base=base)

    # -----------------------------
    # Batch API
    # -----------------------------
    def upload_jsonl(self, path: str) -> dict:
        return File(self.cli).upload(path, purpose="batch", mime_type="application/jsonl")

    def create(
        self,
        input_file_id: str,
        endpoint="/v1/chat/completions",
        window="24h",
        metadata=None,
    ) -> dict:
        body = {
            "input_file_id": input_file_id,
            "endpoint": endpoint,
            "completion_window": window,
            "metadata": metadata or {},
        }
        return self.cli.post("/v1/batches", body).json()

    def retrieve(self, batch_id: str) -> dict:
        return self.cli.get(f"/v1/batches/{batch_id}").json()

    def list(self, limit=20, after=None) -> dict:
        params = {"limit": limit}
        if after:
            params["after"] = after
        return self.cli.get("/v1/batches", **params).json()

    def cancel(self, batch_id: str) -> dict:
        return self.cli.post(f"/v1/batches/{batch_id}/cancel", {}).json()

    def wait(self, batch_id: str, interval=30, timeout=7200) -> dict:
        start = time.time()
        while True:
            data = self.retrieve(batch_id)
            status = data.get("status")
            if self.is_terminal_status(status):
                return data
            if time.time() - start > timeout:
                raise LayBotError(f"batch wait timeout: {batch_id}")
            time.sleep(interval)

    # -----------------------------
    # File download helpers
    # -----------------------------
    def download_output(self, file_id: str) -> str:
        return File(self.cli).download(file_id)

    def download_error(self, file_id: str) -> str:
        return File(self.cli).download(file_id)

    def download_output_to(self, file_id: str, save_path: str) -> str:
        return File(self.cli).download_to(file_id, save_path)

    def download_error_to(self, file_id: str, save_path: str) -> str:
        return File(self.cli).download_to(file_id, save_path)

    def download_and_parse_output(self, file_id: str) -> list[dict]:
        text = self.download_output(file_id)
        return self.parse_batch_output(text)

    def download_and_parse_error(self, file_id: str) -> list[Any]:
        text = self.download_error(file_id)
        return self.parse_jsonl(text)

    # -----------------------------
    # JSONL helpers (re-export)
    # -----------------------------
    @staticmethod
    def parse_jsonl(text: str) -> list[Any]:
        return parse_jsonl(text)

    @staticmethod
    def iter_jsonl(text: str, *, with_line_no: bool = False):
        return iter_jsonl(text, with_line_no=with_line_no)

    @staticmethod
    def iter_jsonl_safe(text: str, *, with_line_no: bool = False):
        return iter_jsonl_safe(text, with_line_no=with_line_no)

    @staticmethod
    def parse_jsonl_file(path: str, *, encoding="utf-8") -> list[Any]:
        return parse_jsonl_file(path, encoding=encoding)

    @staticmethod
    def iter_jsonl_file(path: str, *, with_line_no: bool = False, encoding="utf-8"):
        return iter_jsonl_file(path, with_line_no=with_line_no, encoding=encoding)

    @staticmethod
    def iter_jsonl_file_safe(path: str, *, with_line_no: bool = False, encoding="utf-8"):
        return iter_jsonl_file_safe(path, with_line_no=with_line_no, encoding=encoding)

    @staticmethod
    def dump_jsonl(rows, path: str, *, encoding="utf-8", ensure_ascii=False) -> str:
        return dump_jsonl(rows, path, encoding=encoding, ensure_ascii=ensure_ascii)

    # -----------------------------
    # Batch result helpers
    # -----------------------------
    @classmethod
    def is_terminal_status(cls, status: str | None) -> bool:
        return (status or "").lower() in cls.TERMINAL_STATUSES

    @staticmethod
    def is_success_row(row: dict) -> bool:
        response = row.get("response") or {}
        status_code = response.get("status_code")
        return isinstance(status_code, int) and 200 <= status_code < 300

    @staticmethod
    def extract_custom_id(row: dict) -> str | None:
        v = row.get("custom_id")
        return str(v) if v is not None else None

    @staticmethod
    def extract_response(row: dict) -> dict | None:
        response = row.get("response")
        return response if isinstance(response, dict) else None

    @staticmethod
    def extract_response_body(row: dict) -> dict | None:
        response = row.get("response") or {}
        body = response.get("body")
        return body if isinstance(body, dict) else None

    @staticmethod
    def extract_error(row: dict) -> Any:
        if row.get("error") is not None:
            return row.get("error")
        response = row.get("response") or {}
        body = response.get("body") or {}
        if isinstance(body, dict) and "error" in body:
            return body.get("error")
        return None

    @staticmethod
    def extract_output_text(row: dict) -> str | None:
        """
        尝试从 OpenAI chat.completions batch 结果中提取 assistant 文本。
        兼容:
        - message.content 为字符串
        - message.content 为数组结构
        """
        body = Batch.extract_response_body(row) or {}
        choices = body.get("choices") or []
        if not choices:
            return None

        first = choices[0] if isinstance(choices[0], dict) else {}
        msg = first.get("message") or {}
        content = msg.get("content")

        if isinstance(content, str):
            return content

        if isinstance(content, list):
            texts = []
            for part in content:
                if not isinstance(part, dict):
                    continue
                if part.get("type") == "text" and isinstance(part.get("text"), str):
                    texts.append(part["text"])
            if texts:
                return "".join(texts)

        return None

    @classmethod
    def parse_batch_output(cls, text: str) -> list[dict]:
        """
        将 batch output jsonl 归一化成更易消费的结构。
        """
        rows = []
        for row in cls.iter_jsonl(text):
            rows.append({
                "custom_id": cls.extract_custom_id(row),
                "success": cls.is_success_row(row),
                "response": cls.extract_response(row),
                "body": cls.extract_response_body(row),
                "error": cls.extract_error(row),
                "output_text": cls.extract_output_text(row),
                "raw": row,
            })
        return rows
