########################################################################################################################
# IMPORTS

import base64
import contextlib
import json
import logging
import tempfile
import time
import uuid
from abc import ABC, abstractmethod
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union

from pydantic import BaseModel, ValidationError
from tenacity import retry, retry_if_not_exception_type, stop_after_attempt, wait_exponential

########################################################################################################################
# PARAMETERS

logger = logging.getLogger(__name__)

_LLM_RETRY = retry(
    reraise=True,
    stop=stop_after_attempt(3),
    wait=wait_exponential(multiplier=0.5, min=0.5, max=5),
    retry=retry_if_not_exception_type(ValidationError),
)

########################################################################################################################
# BASE CLASSES


class BaseLLMProvider(ABC):
    """
    Abstract base class for LLM providers.
    """

    @abstractmethod
    def invoke_text(
        self,
        *,
        instructions: str,
        input_texts: Union[str, List[str]],
        model: str,
        image_path: Optional[str] = None,
        image_url: Optional[str] = None,
        reasoning_effort: Optional[str] = None,
    ) -> Tuple[str, Dict[str, Any], Any]:
        """
        Generate text output from the LLM.
        Supports both local images (image_path) and image URLs (image_url).
        """
        pass

    @abstractmethod
    def invoke_structured(
        self,
        *,
        instructions: str,
        input_texts: Union[str, List[str]],
        model_schema: Type[BaseModel],
        model: str,
        image_path: Optional[str] = None,
        image_url: Optional[str] = None,
        reasoning_effort: Optional[str] = None,
        raw_schema: Optional[Dict[str, Any]] = None,
    ) -> Tuple[BaseModel, Dict[str, Any], Any]:
        """
        Generate structured output from the LLM using a Pydantic schema.
        Supports both local images (image_path) and image URLs (image_url).
        """
        pass

    @abstractmethod
    def invoke_embedding(
        self,
        *,
        input_texts: List[str],
        model: str,
    ) -> Tuple[List[List[float]], Dict[str, Any], Any]:
        """
        Generate embeddings for the given input texts.
        """
        pass

    # ------------------------------------------------------------------
    # Batch methods (abstract)
    # ------------------------------------------------------------------

    @abstractmethod
    def submit_batch(
        self,
        requests: List[Dict[str, Any]],
        metadata: Optional[Dict[str, str]] = None,
        wait: bool = False,
        poll_interval: float = 30.0,
        callback: Optional[Callable[[Dict[str, Any]], None]] = None,
    ) -> Tuple[str, Optional[List[Dict[str, Any]]]]:
        """Submit a batch job."""
        pass

    @abstractmethod
    def get_batch_status(self, batch_id: str) -> Dict[str, Any]:
        """Get the current status of a batch job."""
        pass

    @abstractmethod
    def get_batch_results(self, batch_id: str, output_path: Optional[str] = None) -> List[Dict[str, Any]]:
        """Retrieve results from a completed batch."""
        pass

    @abstractmethod
    def cancel_batch(self, batch_id: str) -> Dict[str, Any]:
        """Cancel a running batch job."""
        pass

    @abstractmethod
    def list_batches(self, limit: int = 20, after: Optional[str] = None) -> List[Dict[str, Any]]:
        """List all batches."""
        pass

    @abstractmethod
    def create_batch_file(
        self,
        requests: List[Dict[str, Any]],
        output_path: Optional[str] = None,
    ) -> str:
        """Create a JSONL batch input file."""
        pass

    @abstractmethod
    def upload_batch_file(self, filepath: str) -> str:
        """Upload a batch input file to the provider."""
        pass

    @abstractmethod
    def create_batch(
        self,
        input_file_id: str,
        completion_window: str = "24h",
        metadata: Optional[Dict[str, str]] = None,
    ) -> Dict[str, Any]:
        """Start a batch job with the uploaded file."""
        pass

    # ------------------------------------------------------------------
    # Shared utilities
    # ------------------------------------------------------------------

    @staticmethod
    def _get_image_media_type(image_path: str, default: str = "not_compatible") -> str:
        """
        Determine the media type based on file extension.

        Supported formats: PNG, JPEG, WEBP, non-animated GIF

        Args:
            image_path: Path to the image file
            default: Default value for unsupported extensions

        Returns:
            Media type string (e.g., "image/jpeg", "image/png")
        """
        extension = Path(image_path).suffix.lower()
        media_types = {
            ".jpg": "image/jpeg",
            ".jpeg": "image/jpeg",
            ".png": "image/png",
            ".gif": "image/gif",
            ".webp": "image/webp",
        }
        return media_types.get(extension, default)

    @staticmethod
    def _is_url(path: str) -> bool:
        """Check if the provided path is a URL."""
        return path.startswith(("http://", "https://"))

    @staticmethod
    def _encode_image(image_path: str) -> str:
        """Encodes a local image file to a base64 string."""
        with open(image_path, "rb") as image_file:
            return base64.b64encode(image_file.read()).decode("utf-8")

    @staticmethod
    def _format_usage(meta: Dict[str, Any]) -> str:
        """Format a usage-metadata dict into a readable, emoji-annotated log line."""
        model = meta.get("model", "?")
        input_t = meta.get("input_tokens", 0)
        output_t = meta.get("output_tokens", 0)
        cached_t = meta.get("cached_tokens", 0)
        uncached_t = meta.get("uncached_tokens", input_t)
        latency = meta.get("latency_sec")

        cached_part = f" (💾 {cached_t:,} cached, {uncached_t:,} fresh)" if cached_t else ""
        latency_part = f" | ⚡ {latency:.2f}s" if latency is not None else ""

        return f"💡 {model} | ⬆️  {input_t:,} in{cached_part} | ⬇️  {output_t:,} out{latency_part}"

    @staticmethod
    def _format_batch_usage(results: List[Dict[str, Any]], batch_id: str) -> str:
        """Aggregate token usage across batch results and format as a log line."""
        total = len(results)
        errors = sum(1 for r in results if r.get("error"))
        ok = total - errors

        usages = [r["usage"] for r in results if r.get("usage")]
        total_in = sum(u["input_tokens"] for u in usages)
        total_out = sum(u["output_tokens"] for u in usages)
        total_cached = sum(u["cached_tokens"] for u in usages)

        cached_part = f" (💾 {total_cached:,} cached)" if total_cached else ""
        error_part = f", {errors} ❌" if errors else ""

        return (
            f"📦 Batch {batch_id} | {ok} ✅{error_part} / {total} total"
            f" | ⬆️  {total_in:,} in{cached_part} | ⬇️  {total_out:,} out"
        )

    @staticmethod
    def _parse_jsonl_content(content: str) -> List[Dict[str, Any]]:
        """
        Parse JSONL (JSON Lines) content into a list of dictionaries.

        Args:
            content: Raw JSONL content string

        Returns:
            List of parsed JSON objects
        """
        results = []
        for line in content.strip().split("\n"):
            if line:
                results.append(json.loads(line))
        return results

    @staticmethod
    def _write_jsonl_file(
        lines: List[Dict[str, Any]],
        output_path: Optional[str] = None,
        prefix: str = "batch_input_",
    ) -> str:
        """
        Write a list of dictionaries to a JSONL file.

        Args:
            lines: List of dictionaries to write as JSONL
            output_path: Path to save the file (optional, creates temp file if not provided)
            prefix: Prefix for temp file name

        Returns:
            Path to the created JSONL file
        """
        if output_path is None:
            with tempfile.NamedTemporaryFile(suffix=".jsonl", prefix=prefix, delete=False) as temp_file:
                output_path = temp_file.name

        with open(output_path, "w", encoding="utf-8") as f:
            for line in lines:
                f.write(json.dumps(line, ensure_ascii=False) + "\n")

        return output_path

    def wait_for_batch(
        self,
        batch_id: str,
        poll_interval: float = 30.0,
        timeout: Optional[float] = None,
        callback: Optional[Callable[[Dict[str, Any]], None]] = None,
    ) -> Dict[str, Any]:
        """
        Wait for a batch job to complete, polling periodically.

        Args:
            batch_id: ID of the batch
            poll_interval: Seconds between status checks (default: 30)
            timeout: Maximum seconds to wait (default: None = no timeout)
            callback: Optional function called with status on each poll

        Returns:
            Final batch status

        Raises:
            TimeoutError: If timeout is reached before completion
            RuntimeError: If batch fails or is cancelled
        """
        start_time = time.time()
        terminal_statuses = {"completed", "failed", "expired", "cancelled"}

        while True:
            status = self.get_batch_status(batch_id)

            if callback:
                callback(status)

            current_status = status["status"]
            if current_status in terminal_statuses:
                if current_status == "completed":
                    logger.info(f"Batch {batch_id} completed successfully")
                elif current_status == "failed":
                    raise RuntimeError(f"Batch {batch_id} failed")
                elif current_status == "expired":
                    raise RuntimeError(f"Batch {batch_id} expired before completion")
                elif current_status == "cancelled":
                    raise RuntimeError(f"Batch {batch_id} was cancelled")

                return status

            if timeout and (time.time() - start_time) > timeout:
                raise TimeoutError(f"Timeout waiting for batch {batch_id}")

            request_counts = status.get("request_counts") or {}
            logger.debug(
                f"Batch {batch_id} status: {current_status} "
                f"({request_counts.get('completed', '?')}/{request_counts.get('total', '?')} completed)"
            )

            time.sleep(poll_interval)


class OpenAIProvider(BaseLLMProvider):
    """
    OpenAI provider implementation using the 'client.responses' pattern.
    Supports OpenAI and compatible providers (Ollama, LM Studio, etc.) via base_url.
    """

    def __init__(self, api_key: str, default_model: str = "gpt-5-nano", base_url: Optional[str] = None) -> None:
        from openai import OpenAI

        kwargs = {"api_key": api_key}
        if base_url is not None:
            kwargs["base_url"] = base_url
        self.client = OpenAI(**kwargs)
        self.default_model = default_model

    @staticmethod
    def _batch_to_dict(batch) -> Dict[str, Any]:
        """Convert an OpenAI Batch object to a dictionary."""
        return batch.model_dump()

    def _build_vision_content(
        self,
        input_texts: List[str],
        image_path: Optional[str] = None,
        image_url: Optional[str] = None,
        detail: str = "auto",
    ) -> List[Dict[str, Any]]:
        """
        Builds the content list for the user message following OpenAI Responses API pattern.
        Handles text and optional image (from local file or URL) for Vision capabilities.
        """
        content = []

        # Add text inputs using input_text type
        for text in input_texts:
            content.append({"type": "input_text", "text": text})

        # Add image if provided (prioritize image_url over image_path)
        if image_url:
            content.append({"type": "input_image", "image_url": image_url})
        elif image_path:
            if self._is_url(image_path):
                content.append({"type": "input_image", "image_url": image_path})
            else:
                base64_image = self._encode_image(image_path)
                media_type = self._get_image_media_type(image_path)
                if media_type == "not_compatible":
                    raise ValueError(f"Unsupported image format for file: {image_path}")
                content.append(
                    {"type": "input_image", "image_url": f"data:{media_type};base64,{base64_image}", "detail": detail}
                )

        return content

    @staticmethod
    def _extract_usage_meta(resp) -> Dict[str, Any]:
        """
        Extract usage metadata.
        Note: The structure of 'resp' depends on the specific API version.
        We attempt to read standard usage fields.
        """
        usage = getattr(resp, "usage", None)
        input_tokens = getattr(usage, "input_tokens", 0) if usage else 0
        output_tokens = getattr(usage, "output_tokens", 0) if usage else 0

        # Calculate number of cached tokens
        cached_tokens = 0
        details = getattr(usage, "input_tokens_details", None)
        if details is not None and hasattr(details, "cached_tokens"):
            cached_tokens = details.cached_tokens

        return {
            "response_id": getattr(resp, "id", None),
            "model": getattr(resp, "model", None),
            "input_tokens": input_tokens,
            "cached_tokens": cached_tokens,
            "uncached_tokens": (input_tokens - cached_tokens),
            "output_tokens": output_tokens,
        }

    @_LLM_RETRY
    def invoke_text(
        self,
        *,
        instructions: str,
        input_texts: Union[str, List[str]],
        model: str,
        image_path: Optional[str] = None,
        image_url: Optional[str] = None,
        reasoning_effort: Optional[str] = None,
    ) -> Tuple[str, Dict[str, Any], Any]:
        """
        Generate text output using 'client.responses.create'.

        Args:
            instructions: System instructions for the model
            input_texts: Text input(s) — a single string or list of strings
            model: Model identifier
            image_path: Path to local image file OR image URL (auto-detected)
            image_url: Direct image URL (takes priority over image_path)
            reasoning_effort: Reasoning effort level ("low", "medium", "high")

        Returns:
            Tuple of (output_text, usage_metadata, raw_response)
        """
        if isinstance(input_texts, str):
            input_texts = [input_texts]
        t0 = time.time()

        # Build content with optional image support
        user_content = self._build_vision_content(input_texts, image_path, image_url, detail="high")

        # Construct input list (System + User)
        input_payload = [
            {"role": "system", "content": instructions},
            {"role": "user", "content": user_content},
        ]

        logger.info(f"Prompting {model} (text via responses.create)...")

        # Translate reasoning_effort to OpenAI's reasoning dict
        api_kwargs: Dict[str, Any] = {}
        if reasoning_effort:
            api_kwargs["reasoning"] = {"effort": reasoning_effort}

        resp = self.client.responses.create(
            model=model,
            input=input_payload,
            **api_kwargs,
        )
        # Extract output text
        output_text = getattr(resp, "output_text", "")
        if not output_text and hasattr(resp, "choices"):
            output_text = resp.choices[0].message.content

        meta = self._extract_usage_meta(resp) | {"latency_sec": round(time.time() - t0, 3)}
        logger.info(self._format_usage(meta))
        return output_text, meta, resp

    @_LLM_RETRY
    def invoke_structured(
        self,
        *,
        instructions: str,
        input_texts: Union[str, List[str]],
        model_schema: Type[BaseModel],
        model: str,
        image_path: Optional[str] = None,
        image_url: Optional[str] = None,
        reasoning_effort: Optional[str] = None,
        raw_schema: Optional[Dict[str, Any]] = None,
    ) -> Tuple[BaseModel, Dict[str, Any], Any]:
        """
        Generate structured output using 'client.responses.parse'.

        Args:
            instructions: System instructions for the model
            input_texts: Text input(s) — a single string or list of strings
            model_schema: Pydantic model class for structured output
            model: Model identifier
            image_path: Path to local image file OR image URL (auto-detected)
            image_url: Direct image URL (takes priority over image_path)
            reasoning_effort: Reasoning effort level ("low", "medium", "high")
            raw_schema: Raw JSON schema dict to use instead of Pydantic-derived schema

        Returns:
            Tuple of (pydantic_instance, usage_metadata, raw_response)
        """
        if isinstance(input_texts, str):
            input_texts = [input_texts]
        t0 = time.time()

        # Prepare content (handles text + optional image from path or URL)
        user_content_list = self._build_vision_content(input_texts, image_path, image_url, detail="high")

        # Construct input payload
        input_payload = [
            {"role": "system", "content": instructions},
            {"role": "user", "content": user_content_list},
        ]

        # Common API kwargs
        api_kwargs: Dict[str, Any] = {}
        if reasoning_effort:
            api_kwargs["reasoning"] = {"effort": reasoning_effort}

        if raw_schema is not None:
            text_param = {
                "format": {
                    "type": "json_schema",
                    "name": raw_schema.get("name", "response_schema"),
                    "schema": raw_schema.get("schema", raw_schema),
                }
            }
            logger.info(f"Prompting {model} (structured via responses.create with raw schema)...")
            resp = self.client.responses.create(model=model, input=input_payload, text=text_param, **api_kwargs)
            parsed_output = model_schema.model_validate_json(resp.output_text)
        else:
            logger.info(f"Prompting {model} (structured via responses.parse)...")
            resp = self.client.responses.parse(model=model, input=input_payload, text_format=model_schema, **api_kwargs)
            parsed_output = resp.output_parsed

        meta = self._extract_usage_meta(resp) | {"latency_sec": round(time.time() - t0, 3)}
        logger.info(self._format_usage(meta))
        return parsed_output, meta, resp

    @_LLM_RETRY
    def invoke_embedding(
        self,
        *,
        input_texts: List[str],
        model: str,
    ) -> Tuple[List[List[float]], Dict[str, Any], Any]:
        """
        Generate embeddings using 'client.embeddings.create'.

        Args:
            input_texts: List of text inputs
            model: Model identifier

        Returns:
            Tuple of (embeddings, usage_metadata, raw_response)
        """
        t0 = time.time()

        resp = self.client.embeddings.create(input=input_texts, model=model)
        embeddings = [item.embedding for item in resp.data]

        meta = {
            "model": resp.model,
            "input_tokens": getattr(resp.usage, "prompt_tokens", 0),
            "latency_sec": round(time.time() - t0, 3),
        }
        return embeddings, meta, resp

    # ==========================================================================================================
    # BATCH PROCESSING METHODS
    # ==========================================================================================================

    def _build_batch_request(
        self,
        custom_id: str,
        instructions: str,
        input_texts: List[str],
        model: str,
        image_path: Optional[str] = None,
        image_url: Optional[str] = None,
        reasoning_effort: Optional[str] = None,
        raw_schema: Optional[Dict[str, Any]] = None,
    ) -> Dict[str, Any]:
        """Build a single batch request in the format required by OpenAI Batch API."""
        user_content = self._build_vision_content(input_texts, image_path, image_url, detail="high")

        input_payload = [
            {"role": "system", "content": instructions},
            {"role": "user", "content": user_content},
        ]

        body: Dict[str, Any] = {
            "model": model,
            "input": input_payload,
        }

        if reasoning_effort:
            body["reasoning"] = {"effort": reasoning_effort}

        if raw_schema is not None:
            body["text"] = {
                "format": {
                    "type": "json_schema",
                    "name": raw_schema.get("name", "response_schema"),
                    "schema": raw_schema.get("schema", raw_schema),
                }
            }

        return {
            "custom_id": custom_id,
            "method": "POST",
            "url": "/v1/responses",
            "body": body,
        }

    def create_batch_file(
        self,
        requests: List[Dict[str, Any]],
        output_path: Optional[str] = None,
    ) -> str:
        """
        Create a JSONL batch input file from a list of request configurations.

        Args:
            requests: List of request configurations. Each dict should have:
                - custom_id: Unique identifier
                - instructions: System instructions
                - input_texts: List of input texts
                - model: Model to use (optional, uses default)
                - image_url: Image URL (optional)
                - image_path: Local image path (optional)
                - reasoning_effort: Reasoning effort level (optional)
                - raw_schema: Raw JSON schema dict (optional)
            output_path: Path to save the JSONL file (optional, creates temp file if not provided)

        Returns:
            Path to the created JSONL file
        """
        lines = []
        for req in requests:
            batch_request = self._build_batch_request(
                custom_id=req.get("custom_id", str(uuid.uuid4())),
                instructions=req.get("instructions", ""),
                input_texts=req.get("input_texts", []),
                model=req.get("model", self.default_model),
                image_path=req.get("image_path"),
                image_url=req.get("image_url"),
                reasoning_effort=req.get("reasoning_effort"),
                raw_schema=req.get("raw_schema"),
            )
            lines.append(batch_request)

        filepath = self._write_jsonl_file(lines, output_path, prefix="batch_input_")
        logger.info(f"Created batch file with {len(requests)} requests: {filepath}")
        return filepath

    def upload_batch_file(self, filepath: str) -> str:
        """
        Upload a batch input file to OpenAI.

        Args:
            filepath: Path to the JSONL batch file

        Returns:
            File ID from OpenAI
        """
        with open(filepath, "rb") as f:
            file_obj = self.client.files.create(file=f, purpose="batch")

        logger.info(f"Uploaded batch file: {file_obj.id}")
        return file_obj.id

    def create_batch(
        self,
        input_file_id: str,
        completion_window: str = "24h",
        metadata: Optional[Dict[str, str]] = None,
    ) -> Dict[str, Any]:
        """
        Create a new batch processing job.

        Args:
            input_file_id: ID of the uploaded batch input file
            completion_window: Time window for completion ("24h")
            metadata: Optional metadata dict for the batch

        Returns:
            Batch object with job details
        """
        batch = self.client.batches.create(
            input_file_id=input_file_id,
            endpoint="/v1/responses",
            completion_window=completion_window,
            metadata=metadata,
        )

        batch_dict = self._batch_to_dict(batch)

        logger.info(f"Created batch: {batch.id} (status: {batch.status})")
        return batch_dict

    def get_batch_status(self, batch_id: str) -> Dict[str, Any]:
        """
        Get the current status of a batch job.

        Status values:
            - validating: Input file being validated
            - failed: Validation failed
            - in_progress: Batch is running
            - finalizing: Completing and preparing results
            - completed: Done, results ready
            - expired: Did not complete in time
            - cancelling: Being cancelled
            - cancelled: Was cancelled
        """
        batch = self.client.batches.retrieve(batch_id)

        return self._batch_to_dict(batch)

    def _save_batch_results_to_file(self, content: str, output_path: str) -> None:
        """Save batch results content to a file."""
        with open(output_path, "w", encoding="utf-8") as f:
            f.write(content)
        logger.info(f"Saved batch results to: {output_path}")

    def _fetch_and_parse_file(self, file_id: str) -> List[Dict[str, Any]]:
        """Fetch a file from OpenAI and parse its JSONL content."""
        file_response = self.client.files.content(file_id)
        return self._parse_jsonl_content(file_response.text)

    def _process_output_file(
        self,
        batch_id: str,
        output_file_id: Optional[str],
        output_path: Optional[str],
    ) -> List[Dict[str, Any]]:
        """Process successful batch results from the output file."""
        if not output_file_id:
            if output_path:
                logger.warning(f"No output file for batch {batch_id}, skipping save to {output_path}")
            return []

        file_response = self.client.files.content(output_file_id)
        content = file_response.text

        if output_path:
            self._save_batch_results_to_file(content, output_path)

        return self._parse_jsonl_content(content)

    def _process_error_file(self, batch_id: str, error_file_id: Optional[str]) -> List[Dict[str, Any]]:
        """Process error results from the error file."""
        if not error_file_id:
            return []

        logger.info(f"Batch {batch_id} has errors (file: {error_file_id}). Fetching...")
        try:
            return self._fetch_and_parse_file(error_file_id)
        except Exception as e:
            logger.error(f"Failed to retrieve error file {error_file_id}: {e}")
            return []

    @staticmethod
    def _extract_body_usage(body: Dict[str, Any]) -> Optional[Dict[str, Any]]:
        """
        Extract and flatten token usage from a batch result body dict.

        Returns a dict with input_tokens, output_tokens, cached_tokens, uncached_tokens,
        or None if the body contains no usage field.
        """
        usage = body.get("usage")
        if not usage or not isinstance(usage, dict):
            return None

        input_tokens = usage.get("input_tokens", 0)
        output_tokens = usage.get("output_tokens", 0)

        cached_tokens = 0
        details = usage.get("input_tokens_details")
        if isinstance(details, dict):
            cached_tokens = details.get("cached_tokens", 0)

        return {
            "input_tokens": input_tokens,
            "output_tokens": output_tokens,
            "cached_tokens": cached_tokens,
            "uncached_tokens": input_tokens - cached_tokens,
        }

    @staticmethod
    def _extract_output_text_from_responses_output(output_items: Any) -> Optional[str]:
        if not isinstance(output_items, list):
            return None

        for item in output_items:
            if not isinstance(item, dict) or item.get("type") != "message":
                continue

            content_blocks = item.get("content", [])
            if not isinstance(content_blocks, list):
                continue

            for block in content_blocks:
                if isinstance(block, dict) and block.get("type") == "output_text":
                    return block.get("text")

        return None

    @staticmethod
    def _extract_output_text_from_body(body: Dict[str, Any]) -> Optional[str]:
        output_text = body.get("output_text")
        if output_text is not None:
            return output_text

        choices = body.get("choices", [])
        if choices:
            choice_output_text = (choices[0].get("message") or {}).get("content", "")
            if choice_output_text is not None:
                return choice_output_text

        return OpenAIProvider._extract_output_text_from_responses_output(body.get("output", []))

    @staticmethod
    def _flatten_batch_result(raw_result: Dict[str, Any]) -> Dict[str, Any]:
        """Flatten a raw batch result into the standardized BatchResult format."""
        custom_id = raw_result.get("custom_id", "")
        error = raw_result.get("error")
        response = raw_result.get("response")

        output_text = None
        usage = None
        if isinstance(response, dict):
            body = response.get("body", {})
            if isinstance(body, dict):
                output_text = OpenAIProvider._extract_output_text_from_body(body)
                usage = OpenAIProvider._extract_body_usage(body)
                if not error and body.get("error"):
                    error = body["error"]

        return {
            "custom_id": custom_id,
            "output_text": output_text,
            "error": error,
            "usage": usage,
        }

    def get_batch_results(
        self,
        batch_id: str,
        output_path: Optional[str] = None,
    ) -> List[Dict[str, Any]]:
        """
        Retrieve results from a completed batch.

        Args:
            batch_id: ID of the batch
            output_path: Optional path to save the raw JSONL output

        Returns:
            List of result dictionaries, each containing:
                - custom_id: The original request ID
                - output_text: The generated text (flattened)
                - error: Error data (if failed)
        """
        status = self.get_batch_status(batch_id)

        if status["status"] != "completed":
            raise RuntimeError(f"Batch {batch_id} is not completed (status: {status['status']})")

        # Process both output and error files
        output_results = self._process_output_file(
            batch_id,
            status.get("output_file_id"),
            output_path,
        )
        error_results = self._process_error_file(batch_id, status.get("error_file_id"))

        raw_results = output_results + error_results

        # Log warning if batch has no results
        if not raw_results and not status.get("output_file_id") and not status.get("error_file_id"):
            logger.warning(f"Batch {batch_id} completed but has no output or error file.")

        results = [self._flatten_batch_result(r) for r in raw_results]

        logger.info(f"Retrieved {len(results)} results from batch {batch_id}")
        logger.info(self._format_batch_usage(results, batch_id))
        return results

    def get_partial_batch_results(self, batch_id: str) -> List[Dict[str, Any]]:
        """
        Retrieve available results from a batch regardless of completion status.

        Unlike get_batch_results, this does not raise on non-completed batches,
        making it suitable for expired batches with partial output.

        Returns:
            List of flattened result dicts (custom_id, output_text, error).
        """
        status = self.get_batch_status(batch_id)
        output_results = self._process_output_file(batch_id, status.get("output_file_id"), None)
        error_results = self._process_error_file(batch_id, status.get("error_file_id"))
        raw_results = output_results + error_results
        results = [self._flatten_batch_result(r) for r in raw_results]
        logger.info(self._format_batch_usage(results, batch_id))
        return results

    def get_batch_errors(self, batch_id: str) -> List[Dict[str, Any]]:
        """Retrieve errors from a batch (if any)."""
        status = self.get_batch_status(batch_id)

        error_file_id = status["error_file_id"]
        if not error_file_id:
            return []

        file_response = self.client.files.content(error_file_id)
        content = file_response.text

        errors = []
        for line in content.strip().split("\n"):
            if line:
                errors.append(json.loads(line))

        return errors

    def cancel_batch(self, batch_id: str) -> Dict[str, Any]:
        """Cancel a running batch job."""
        batch = self.client.batches.cancel(batch_id)
        logger.info(f"Cancelling batch {batch_id}")

        return self._batch_to_dict(batch)

    def list_batches(self, limit: int = 20, after: Optional[str] = None) -> List[Dict[str, Any]]:
        """List all batches."""
        batches = self.client.batches.list(limit=limit, after=after)

        return [self._batch_to_dict(b) for b in batches.data]

    def submit_batch(
        self,
        requests: List[Dict[str, Any]],
        metadata: Optional[Dict[str, str]] = None,
        wait: bool = False,
        poll_interval: float = 30.0,
        callback: Optional[Callable[[Dict[str, Any]], None]] = None,
    ) -> Tuple[str, Optional[List[Dict[str, Any]]]]:
        """
        High-level method to submit a batch job (create file, upload, and start batch).

        Args:
            requests: List of request configurations (see create_batch_file for format)
            metadata: Optional metadata for the batch
            wait: If True, wait for completion and return results
            poll_interval: Seconds between status checks when waiting
            callback: Optional callback for status updates when waiting

        Returns:
            Tuple of (batch_id, results). Results is None if wait=False.
        """
        # Step 1: Create the batch file
        filepath = self.create_batch_file(requests)

        try:
            # Step 2: Upload the file
            file_id = self.upload_batch_file(filepath)

            # Step 3: Create the batch
            batch = self.create_batch(
                input_file_id=file_id,
                metadata=metadata,
            )

            batch_id = batch["id"]

            if not wait:
                return batch_id, None

            # Step 4: Wait for completion
            self.wait_for_batch(batch_id, poll_interval=poll_interval, callback=callback)

            # Step 5: Get results
            results = self.get_batch_results(batch_id)

            return batch_id, results

        finally:
            # Clean up temp file
            Path(filepath).unlink(missing_ok=True)


########################################################################################################################
# GEMINI PROVIDER


class GeminiProvider(BaseLLMProvider):
    """
    Google Gemini provider using the OpenAI compatibility layer.

    Interactive calls use ``client.chat.completions`` via the OpenAI SDK pointed at
    Gemini's compatibility endpoint.  Batch file upload/download uses the native
    ``google-genai`` SDK (lazy-loaded) since the compatibility layer does not support
    file operations.
    """

    GEMINI_BASE_URL = "https://generativelanguage.googleapis.com/v1beta/openai/"

    def __init__(self, api_key: str, default_model: str = "gemini-3.1-flash-lite") -> None:
        from openai import OpenAI

        self.client = OpenAI(api_key=api_key, base_url=self.GEMINI_BASE_URL)
        self.default_model = default_model
        self._api_key = api_key
        self._genai_client = None
        self._genai_types = None
        self._native_batch_ids: set[str] = set()
        self._native_batch_totals: Dict[str, int] = {}

    def _get_genai_client(self):
        """Lazy-load the native genai client (used only for batch file upload/download)."""
        if self._genai_client is None:
            from google import genai

            self._genai_client = genai.Client(api_key=self._api_key)
            self._genai_types = genai.types
        return self._genai_client

    # ------------------------------------------------------------------
    # Vision content (Chat Completions format)
    # ------------------------------------------------------------------

    def _build_chat_content(
        self,
        input_texts: List[str],
        image_path: Optional[str] = None,
        image_url: Optional[str] = None,
    ) -> List[Dict[str, Any]]:
        """Build content blocks in OpenAI Chat Completions format for Gemini compatibility."""
        content: List[Dict[str, Any]] = []

        for text in input_texts:
            content.append({"type": "text", "text": text})

        if image_url:
            content.append({"type": "image_url", "image_url": {"url": image_url}})
        elif image_path:
            if self._is_url(image_path):
                content.append({"type": "image_url", "image_url": {"url": image_path}})
            else:
                base64_image = self._encode_image(image_path)
                media_type = self._get_image_media_type(image_path, default="image/jpeg")
                content.append({"type": "image_url", "image_url": {"url": f"data:{media_type};base64,{base64_image}"}})

        return content

    @staticmethod
    def _extract_usage_meta(resp) -> Dict[str, Any]:
        """Extract usage metadata from Chat Completions response."""
        usage = getattr(resp, "usage", None)
        input_tokens = getattr(usage, "prompt_tokens", 0) if usage else 0
        output_tokens = getattr(usage, "completion_tokens", 0) if usage else 0

        # Check for cached tokens from implicit caching
        cached_tokens = 0
        details = getattr(usage, "prompt_tokens_details", None)
        if details is not None:
            raw = getattr(details, "cached_tokens", None)
            if isinstance(raw, int):
                cached_tokens = raw

        return {
            "response_id": getattr(resp, "id", None),
            "model": getattr(resp, "model", None),
            "input_tokens": input_tokens,
            "cached_tokens": cached_tokens,
            "uncached_tokens": input_tokens - cached_tokens,
            "output_tokens": output_tokens,
        }

    # ------------------------------------------------------------------
    # Core methods
    # ------------------------------------------------------------------

    @_LLM_RETRY
    def invoke_text(
        self,
        *,
        instructions: str,
        input_texts: Union[str, List[str]],
        model: str,
        image_path: Optional[str] = None,
        image_url: Optional[str] = None,
        reasoning_effort: Optional[str] = None,
    ) -> Tuple[str, Dict[str, Any], Any]:
        """Generate text output via Gemini OpenAI compatibility layer."""
        if isinstance(input_texts, str):
            input_texts = [input_texts]
        t0 = time.time()

        user_content = self._build_chat_content(input_texts, image_path, image_url)

        messages = [
            {"role": "system", "content": instructions},
            {"role": "user", "content": user_content},
        ]

        api_kwargs: Dict[str, Any] = {}
        if reasoning_effort:
            api_kwargs["reasoning_effort"] = reasoning_effort

        logger.info(f"Prompting {model} (text via Gemini chat.completions.create)...")

        resp = self.client.chat.completions.create(model=model, messages=messages, **api_kwargs)

        output_text = resp.choices[0].message.content or ""
        meta = self._extract_usage_meta(resp) | {"latency_sec": round(time.time() - t0, 3)}
        logger.info(self._format_usage(meta))
        return output_text, meta, resp

    @_LLM_RETRY
    def invoke_structured(
        self,
        *,
        instructions: str,
        input_texts: Union[str, List[str]],
        model_schema: Type[BaseModel],
        model: str,
        image_path: Optional[str] = None,
        image_url: Optional[str] = None,
        reasoning_effort: Optional[str] = None,
        raw_schema: Optional[Dict[str, Any]] = None,
    ) -> Tuple[BaseModel, Dict[str, Any], Any]:
        """Generate structured output via Gemini OpenAI compatibility layer."""
        if isinstance(input_texts, str):
            input_texts = [input_texts]
        t0 = time.time()

        user_content = self._build_chat_content(input_texts, image_path, image_url)

        messages = [
            {"role": "system", "content": instructions},
            {"role": "user", "content": user_content},
        ]

        api_kwargs: Dict[str, Any] = {}
        if reasoning_effort:
            api_kwargs["reasoning_effort"] = reasoning_effort

        if raw_schema is not None:
            response_format = {
                "type": "json_schema",
                "json_schema": {
                    "name": raw_schema.get("name", "response_schema"),
                    "schema": raw_schema.get("schema", raw_schema),
                },
            }

            logger.info(f"Prompting {model} (structured via Gemini chat.completions.create with raw schema)...")
            resp = self.client.chat.completions.create(
                model=model,
                messages=messages,
                response_format=response_format,
                **api_kwargs,
            )
            parsed_output = model_schema.model_validate_json(resp.choices[0].message.content)
        else:
            logger.info(f"Prompting {model} (structured via Gemini beta.chat.completions.parse)...")
            resp = self.client.beta.chat.completions.parse(
                model=model,
                messages=messages,
                response_format=model_schema,
                **api_kwargs,
            )
            parsed_output = resp.choices[0].message.parsed

        meta = self._extract_usage_meta(resp) | {"latency_sec": round(time.time() - t0, 3)}
        logger.info(self._format_usage(meta))
        return parsed_output, meta, resp

    @_LLM_RETRY
    def invoke_embedding(
        self,
        *,
        input_texts: List[str],
        model: str,
    ) -> Tuple[List[List[float]], Dict[str, Any], Any]:
        """
        Generate embeddings via Gemini OpenAI compatibility layer.

        Args:
            input_texts: List of text inputs
            model: Model identifier

        Returns:
            Tuple of (embeddings, usage_metadata, raw_response)
        """
        t0 = time.time()

        resp = self.client.embeddings.create(input=input_texts, model=model)
        embeddings = [item.embedding for item in resp.data]

        meta = {
            "model": resp.model,
            "input_tokens": getattr(resp.usage, "prompt_tokens", 0),
            "latency_sec": round(time.time() - t0, 3),
        }
        return embeddings, meta, resp

    # ------------------------------------------------------------------
    # Batch processing (via OpenAI compatibility + native genai for files)
    # ------------------------------------------------------------------

    def _download_file(self, file_id: str) -> str:
        """Download a file via the native genai client."""
        genai_client = self._get_genai_client()
        file_bytes = genai_client.files.download(file=file_id)
        return file_bytes.decode("utf-8") if isinstance(file_bytes, (bytes, bytearray)) else str(file_bytes)

    def _build_native_inlined_requests(self, requests: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """Build native Gemini inlined batch requests with structured-output config."""
        inlined_requests: List[Dict[str, Any]] = []

        for req in requests:
            custom_id = req.get("custom_id", str(uuid.uuid4()))
            instructions = req.get("instructions", "")
            input_texts = req.get("input_texts", [])
            image_url = req.get("image_url")
            image_path = req.get("image_path")
            raw_schema = req.get("raw_schema")
            reasoning_effort = req.get("reasoning_effort")

            response_json_schema = None
            if raw_schema and isinstance(raw_schema, dict):
                response_json_schema = raw_schema.get("schema", raw_schema)

            config: Dict[str, Any] = {
                "system_instruction": instructions,
                "response_mime_type": "application/json",
            }
            if response_json_schema:
                config["response_json_schema"] = response_json_schema
            if reasoning_effort:
                config["thinking_config"] = {"thinking_level": reasoning_effort}

            parts: List[Dict[str, Any]] = [{"text": "\n\n".join(input_texts)}]

            image_uri = image_url or image_path
            if image_uri:
                parts.append({"file_data": {"file_uri": image_uri}})

            inlined_requests.append(
                {
                    "contents": [{"role": "user", "parts": parts}],
                    "config": config,
                    "metadata": {"custom_id": str(custom_id)},
                }
            )

        return inlined_requests

    @staticmethod
    def _map_native_job_state(state_value: str) -> str:
        """Map native Gemini job states to OpenAI-like batch statuses used by shared wait loop."""
        state_map = {
            "JOB_STATE_SUCCEEDED": "completed",
            "JOB_STATE_PARTIALLY_SUCCEEDED": "completed",
            "JOB_STATE_FAILED": "failed",
            "JOB_STATE_CANCELLED": "cancelled",
            "JOB_STATE_EXPIRED": "expired",
            "JOB_STATE_QUEUED": "in_progress",
            "JOB_STATE_PENDING": "in_progress",
            "JOB_STATE_RUNNING": "in_progress",
            "JOB_STATE_CANCELLING": "in_progress",
            "JOB_STATE_UPDATING": "in_progress",
            "JOB_STATE_PAUSED": "in_progress",
            "JOB_STATE_UNSPECIFIED": "in_progress",
        }
        return state_map.get(state_value, "in_progress")

    def create_batch_file(
        self,
        requests: List[Dict[str, Any]],
        output_path: Optional[str] = None,
    ) -> str:
        """Create a JSONL batch input file in OpenAI Chat Completions format."""
        lines = []
        for req in requests:
            custom_id = req.get("custom_id", str(uuid.uuid4()))
            instructions = req.get("instructions", "")
            input_texts = req.get("input_texts", [])
            model = req.get("model", self.default_model)
            image_url = req.get("image_url")
            image_path = req.get("image_path")
            reasoning_effort = req.get("reasoning_effort")
            raw_schema = req.get("raw_schema")

            batch_instructions = instructions

            # Translate raw_schema → instructions for batch mode
            # (Gemini batch validator rejects full JSON-schema payloads in batch lines)
            if raw_schema and isinstance(raw_schema, dict):
                schema_obj = raw_schema.get("schema", raw_schema)
                properties = schema_obj.get("properties", {}) if isinstance(schema_obj, dict) else {}
                required = schema_obj.get("required", []) if isinstance(schema_obj, dict) else []

                if properties:
                    keys = list(properties.keys())
                    required_keys = required if isinstance(required, list) and required else keys
                    batch_instructions = (
                        f"{instructions}\n\n"
                        "IMPORTANT: Respond with ONLY a valid JSON object (no markdown, no code fences). "
                        f"Include exactly these keys: {', '.join(keys)}. "
                        f"Required keys: {', '.join(required_keys)}."
                    )

            extra_kwargs: Dict[str, Any] = {}
            if reasoning_effort:
                extra_kwargs["reasoning_effort"] = reasoning_effort

            # Gemini batch validator is stricter than sync chat endpoint for text-only payloads.
            if image_url or image_path:
                user_content = self._build_chat_content(input_texts, image_path, image_url)
            else:
                user_content = "\n\n".join(input_texts)
            messages = [
                {"role": "system", "content": batch_instructions},
                {"role": "user", "content": user_content},
            ]

            lines.append(
                {
                    "custom_id": custom_id,
                    "method": "POST",
                    "url": "/v1/chat/completions",
                    "body": {
                        "model": model,
                        "messages": messages,
                        **extra_kwargs,
                    },
                }
            )

        filepath = self._write_jsonl_file(lines, output_path, prefix="gemini_batch_input_")
        logger.info(f"Created Gemini batch file with {len(requests)} requests: {filepath}")
        return filepath

    def upload_batch_file(self, filepath: str) -> str:
        """Upload a JSONL batch input file via the native genai File API."""
        genai_client = self._get_genai_client()
        upload_config = self._genai_types.UploadFileConfig(
            display_name=Path(filepath).stem,
            mime_type="jsonl",
        )
        uploaded_file = genai_client.files.upload(file=filepath, config=upload_config)
        file_name = getattr(uploaded_file, "name", None)
        if not file_name:
            raise RuntimeError("Gemini file upload succeeded but returned no file name")

        logger.info(f"Uploaded Gemini batch file: {file_name}")
        return file_name

    def create_batch(
        self,
        input_file_id: str,
        completion_window: str = "24h",
        metadata: Optional[Dict[str, str]] = None,
    ) -> Dict[str, Any]:
        """Create a Gemini batch job via the OpenAI compatibility layer."""
        batch = self.client.batches.create(
            input_file_id=input_file_id,
            endpoint="/v1/chat/completions",
            completion_window=completion_window,
            metadata=metadata,
        )

        batch_dict = batch.model_dump()
        logger.info(f"Created Gemini batch: {batch.id} (status: {batch.status})")
        return batch_dict

    def get_batch_status(self, batch_id: str) -> Dict[str, Any]:
        """Get the current status of a Gemini batch job."""
        if batch_id in self._native_batch_ids:
            genai_client = self._get_genai_client()
            job = genai_client.batches.get(name=batch_id)
            state_obj = getattr(job, "state", None)
            state_value = getattr(state_obj, "value", str(state_obj))
            mapped_status = self._map_native_job_state(state_value)

            completion_stats = getattr(job, "completion_stats", None)
            successful = int(getattr(completion_stats, "successful_count", 0) or 0)
            failed = int(getattr(completion_stats, "failed_count", 0) or 0)
            incomplete = int(getattr(completion_stats, "incomplete_count", 0) or 0)

            source = getattr(job, "src", None)
            source_requests = getattr(source, "inlined_requests", None) if source is not None else None
            total = self._native_batch_totals.get(batch_id)
            if total is None:
                total = len(source_requests) if source_requests is not None else (successful + failed + incomplete)
            completed = successful + failed
            if completed == 0 and mapped_status == "completed" and total:
                completed = total

            return {
                "id": batch_id,
                "status": mapped_status,
                "request_counts": {
                    "completed": completed,
                    "total": total,
                },
                "raw_state": state_value,
            }

        batch = self.client.batches.retrieve(batch_id)
        return batch.model_dump()

    @staticmethod
    def _extract_native_response_text(response: Any) -> str:
        """Extract text from a native Gemini inlined response object."""
        output_text = getattr(response, "text", None)
        if output_text is not None:
            return output_text or ""

        candidates = getattr(response, "candidates", None) or []
        if not candidates:
            return ""

        content = getattr(candidates[0], "content", None)
        parts = getattr(content, "parts", None) if content is not None else None
        if not parts:
            return ""

        return "".join(getattr(part, "text", "") or "" for part in parts)

    def _parse_native_inlined_result(self, inlined: Any, index: int) -> Dict[str, Any]:
        """Parse one native Gemini inlined response to BatchResult format."""
        metadata = getattr(inlined, "metadata", None) or {}
        custom_id = metadata.get("custom_id") or str(index)

        error = getattr(inlined, "error", None)
        if error is not None:
            return {
                "custom_id": custom_id,
                "output_text": None,
                "error": {
                    "code": getattr(error, "code", None),
                    "message": getattr(error, "message", str(error)),
                },
                "usage": None,
            }

        response = getattr(inlined, "response", None)
        if response is None:
            return {
                "custom_id": custom_id,
                "output_text": None,
                "error": {"message": "Missing response in native batch result"},
                "usage": None,
            }

        usage = None
        usage_metadata = getattr(response, "usage_metadata", None)
        if usage_metadata:
            input_tokens = getattr(usage_metadata, "prompt_token_count", 0) or 0
            output_tokens = getattr(usage_metadata, "candidates_token_count", 0) or 0
            cached_tokens = getattr(usage_metadata, "cached_content_token_count", 0) or 0
            usage = {
                "input_tokens": input_tokens,
                "output_tokens": output_tokens,
                "cached_tokens": cached_tokens,
                "uncached_tokens": input_tokens - cached_tokens,
            }

        return {
            "custom_id": custom_id,
            "output_text": self._extract_native_response_text(response),
            "error": None,
            "usage": usage,
        }

    def _get_native_batch_results(self, batch_id: str) -> List[Dict[str, Any]]:
        """Retrieve and flatten results from a completed native Gemini batch job."""
        genai_client = self._get_genai_client()
        job = genai_client.batches.get(name=batch_id)

        state_obj = getattr(job, "state", None)
        state_value = getattr(state_obj, "value", str(state_obj))
        if self._map_native_job_state(state_value) != "completed":
            raise RuntimeError(f"Batch {batch_id} is not completed (status: {state_value})")

        destination = getattr(job, "dest", None)
        inlined_responses = getattr(destination, "inlined_responses", None) if destination else None
        if not inlined_responses:
            return []

        results = [
            self._parse_native_inlined_result(inlined, index) for index, inlined in enumerate(inlined_responses, 1)
        ]

        logger.info(f"Retrieved {len(results)} results from native Gemini batch {batch_id}")
        return results

    @staticmethod
    def _extract_body_usage(body: Dict[str, Any]) -> Optional[Dict[str, Any]]:
        """
        Extract and flatten token usage from a Gemini compat batch result body dict.

        Gemini chat-completions format uses prompt_tokens / completion_tokens /
        prompt_tokens_details instead of OpenAI's input_tokens / output_tokens /
        input_tokens_details.  Returns a normalised dict with the same keys as
        OpenAIProvider._extract_body_usage, or None if no usage field is present.
        """
        usage = body.get("usage")
        if not usage or not isinstance(usage, dict):
            return None

        input_tokens = usage.get("prompt_tokens", 0)
        output_tokens = usage.get("completion_tokens", 0)

        cached_tokens = 0
        details = usage.get("prompt_tokens_details")
        if isinstance(details, dict):
            cached_tokens = details.get("cached_tokens", 0)

        return {
            "input_tokens": input_tokens,
            "output_tokens": output_tokens,
            "cached_tokens": cached_tokens,
            "uncached_tokens": input_tokens - cached_tokens,
        }

    @staticmethod
    def _flatten_compat_batch_result(raw_result: Dict[str, Any]) -> Dict[str, Any]:
        """Flatten one OpenAI-compat batch JSONL entry to BatchResult format."""
        custom_id = raw_result.get("custom_id", "")
        error = raw_result.get("error")

        output_text = None
        usage = None
        response = raw_result.get("response")
        if response and isinstance(response, dict) and "body" in response:
            body = response["body"]
            output_text = body.get("output_text")
            if output_text is None:
                choices = body.get("choices", [])
                if choices:
                    output_text = (choices[0].get("message") or {}).get("content", "")
            usage = GeminiProvider._extract_body_usage(body)

        return {
            "custom_id": custom_id,
            "output_text": output_text,
            "error": error,
            "usage": usage,
        }

    def _get_compat_batch_results(self, batch_id: str, output_path: Optional[str] = None) -> List[Dict[str, Any]]:
        """Retrieve and flatten results from OpenAI-compatible Gemini batch jobs."""
        status = self.get_batch_status(batch_id)

        if status["status"] != "completed":
            raise RuntimeError(f"Batch {batch_id} is not completed (status: {status['status']})")

        raw_results: List[Dict[str, Any]] = []

        output_file_id = status.get("output_file_id")
        if output_file_id:
            content = self._download_file(output_file_id)
            if output_path:
                with open(output_path, "w", encoding="utf-8") as f:
                    f.write(content)
                logger.info(f"Saved batch results to: {output_path}")
            raw_results.extend(self._parse_jsonl_content(content))

        error_file_id = status.get("error_file_id")
        if error_file_id:
            try:
                content = self._download_file(error_file_id)
                raw_results.extend(self._parse_jsonl_content(content))
            except Exception as e:
                logger.error(f"Failed to retrieve error file {error_file_id}: {e}")

        results = [self._flatten_compat_batch_result(raw_result) for raw_result in raw_results]
        logger.info(f"Retrieved {len(results)} results from batch {batch_id}")
        logger.info(self._format_batch_usage(results, batch_id))
        return results

    def get_batch_results(
        self,
        batch_id: str,
        output_path: Optional[str] = None,
    ) -> List[Dict[str, Any]]:
        """
        Retrieve results from a completed Gemini batch job.

        Returns flat BatchResult dicts: {"custom_id": ..., "output_text": ..., "error": ...}
        """
        if batch_id in self._native_batch_ids:
            results = self._get_native_batch_results(batch_id)
            logger.info(self._format_batch_usage(results, batch_id))
            return results

        return self._get_compat_batch_results(batch_id, output_path)

    def get_partial_batch_results(self, batch_id: str) -> List[Dict[str, Any]]:
        """
        Retrieve available results from a batch regardless of completion status.

        For native Gemini batches, attempts to get results (returns empty if not ready).
        For compat-layer batches, fetches output + error files without status guard.
        """
        if batch_id in self._native_batch_ids:
            try:
                return self._get_native_batch_results(batch_id)
            except RuntimeError:
                return []

        # Compat path — like _get_compat_batch_results but without status check
        status = self.get_batch_status(batch_id)
        raw_results: List[Dict[str, Any]] = []

        output_file_id = status.get("output_file_id")
        if output_file_id:
            content = self._download_file(output_file_id)
            raw_results.extend(self._parse_jsonl_content(content))

        error_file_id = status.get("error_file_id")
        if error_file_id:
            try:
                content = self._download_file(error_file_id)
                raw_results.extend(self._parse_jsonl_content(content))
            except Exception as e:
                logger.error(f"Failed to retrieve error file {error_file_id}: {e}")

        results = [self._flatten_compat_batch_result(r) for r in raw_results]
        logger.info(self._format_batch_usage(results, batch_id))
        return results

    def get_batch_errors(self, batch_id: str) -> List[Dict[str, Any]]:
        """Return only error entries from batch results."""
        try:
            results = self.get_batch_results(batch_id)
        except RuntimeError:
            return []
        return [r for r in results if r.get("error")]

    def cancel_batch(self, batch_id: str) -> Dict[str, Any]:
        """Cancel a Gemini batch job."""
        if batch_id in self._native_batch_ids:
            genai_client = self._get_genai_client()
            job = genai_client.batches.cancel(name=batch_id)
            return job.model_dump()

        batch = self.client.batches.cancel(batch_id)
        logger.info(f"Cancelling batch {batch_id}")
        return batch.model_dump()

    def list_batches(self, limit: int = 20, after: Optional[str] = None) -> List[Dict[str, Any]]:
        """List Gemini batches."""
        batches = self.client.batches.list(limit=limit, after=after)
        result = []
        for b in batches.data:
            d = b.model_dump()
            if isinstance(d.get("created_at"), str):
                with contextlib.suppress(ValueError, TypeError):
                    d["created_at"] = int(d["created_at"])
            result.append(d)
        return result

    def submit_batch(
        self,
        requests: List[Dict[str, Any]],
        metadata: Optional[Dict[str, str]] = None,
        wait: bool = False,
        poll_interval: float = 30.0,
        callback: Optional[Callable[[Dict[str, Any]], None]] = None,
    ) -> Tuple[str, Optional[List[Dict[str, Any]]]]:
        """
        Submit a batch job via the Gemini OpenAI compatibility layer.

        Creates a JSONL batch file in OpenAI Chat Completions format, uploads it
        via the native genai SDK, and submits the batch via the compatibility layer.
        """
        if not requests:
            raise ValueError("requests must not be empty")

        # Prefer native Gemini inlined batches for structured output, which
        # supports response_mime_type/response_json_schema directly.
        all_have_structured_schema = all(isinstance(req.get("raw_schema"), dict) for req in requests)
        models = {req.get("model", self.default_model) for req in requests}
        if all_have_structured_schema and len(models) == 1:
            native_model = next(iter(models))
            native_requests = self._build_native_inlined_requests(requests)
            genai_client = self._get_genai_client()

            create_config: Dict[str, Any] = {}
            if metadata:
                display_name = metadata.get("task") or metadata.get("crawler") or "structured-output-batch"
                create_config["display_name"] = display_name

            native_job = genai_client.batches.create(
                model=native_model,
                src=native_requests,
                config=create_config or None,
            )
            batch_id = native_job.name
            self._native_batch_ids.add(batch_id)
            self._native_batch_totals[batch_id] = len(native_requests)

            logger.info(f"Created native Gemini structured batch: {batch_id}")

            if not wait:
                return batch_id, None

            self.wait_for_batch(batch_id, poll_interval=poll_interval, callback=callback)
            results = self.get_batch_results(batch_id)
            return batch_id, results

        filepath = self.create_batch_file(requests)

        try:
            file_id = self.upload_batch_file(filepath)

            batch = self.create_batch(
                input_file_id=file_id,
                metadata=metadata,
            )

            batch_id = batch["id"]

            if not wait:
                return batch_id, None

            self.wait_for_batch(batch_id, poll_interval=poll_interval, callback=callback)
            results = self.get_batch_results(batch_id)
            return batch_id, results

        finally:
            Path(filepath).unlink(missing_ok=True)


########################################################################################################################
# MAIN INTERFACE


class LLMInterface:
    """
    Unified interface for working with Large Language Models.
    """

    SUPPORTED_PROVIDERS = {
        "openai": OpenAIProvider,
        "gemini": GeminiProvider,
    }

    def __init__(self, config) -> None:
        if "llm" not in config:
            raise ValueError("Configuration must contain 'llm' section")

        llm_config = config["llm"]
        self.provider_name = llm_config.get("provider", "openai").lower()

        if self.provider_name not in self.SUPPORTED_PROVIDERS:
            raise ValueError(f"Unsupported provider '{self.provider_name}'")

        api_key = llm_config.get("api_key")
        if not api_key:
            raise ValueError(f"API key required for provider '{self.provider_name}'")

        if "model" in llm_config:
            default_model = llm_config["model"]
        elif self.provider_name == "gemini":
            default_model = "gemini-3.1-flash-lite"
        else:
            default_model = "gpt-5-nano"

        provider_class = self.SUPPORTED_PROVIDERS[self.provider_name]
        provider_kwargs: Dict[str, Any] = {"api_key": api_key, "default_model": default_model}
        base_url = llm_config.get("base_url")
        if base_url and self.provider_name == "openai":
            provider_kwargs["base_url"] = base_url
        self.provider: BaseLLMProvider = provider_class(**provider_kwargs)
        self.default_model = default_model

        logger.info(f"Initialized LLM interface with provider: {self.provider_name}")

    def invoke_text(
        self,
        *,
        instructions: str,
        input_texts: Union[str, List[str]],
        image_path: Optional[str] = None,
        image_url: Optional[str] = None,
        model: Optional[str] = None,
        reasoning_effort: Optional[str] = None,
    ) -> Tuple[str, Dict[str, Any], Any]:
        """
        Generate text output with optional Vision support.

        Args:
            instructions: System instructions for the model
            input_texts: Text input(s) — a single string or list of strings
            image_path: Path to local image file OR image URL (auto-detected)
            image_url: Direct image URL (takes priority over image_path)
            model: Model identifier (uses default if not specified)
            reasoning_effort: Reasoning effort level ("low", "medium", "high")

        Returns:
            Tuple of (output_text, usage_metadata, raw_response)
        """
        model = model or self.default_model
        return self.provider.invoke_text(
            instructions=instructions,
            input_texts=input_texts,
            model=model,
            image_path=image_path,
            image_url=image_url,
            reasoning_effort=reasoning_effort,
        )

    def invoke_structured(
        self,
        *,
        instructions: str,
        input_texts: Union[str, List[str]],
        model_schema: Type[BaseModel],
        image_path: Optional[str] = None,
        image_url: Optional[str] = None,
        model: Optional[str] = None,
        reasoning_effort: Optional[str] = None,
        raw_schema: Optional[Dict[str, Any]] = None,
    ) -> Tuple[BaseModel, Dict[str, Any], Any]:
        """
        Generate structured output with Vision support.

        Args:
            instructions: System instructions for the model
            input_texts: Text input(s) — a single string or list of strings
            model_schema: Pydantic model class for structured output
            image_path: Path to local image file OR image URL (auto-detected)
            image_url: Direct image URL (takes priority over image_path)
            model: Model identifier (uses default if not specified)
            reasoning_effort: Reasoning effort level ("low", "medium", "high")
            raw_schema: Raw JSON schema dict to use instead of Pydantic-derived schema

        Returns:
            Tuple of (pydantic_instance, usage_metadata, raw_response)
        """
        model = model or self.default_model
        return self.provider.invoke_structured(
            instructions=instructions,
            input_texts=input_texts,
            model_schema=model_schema,
            model=model,
            image_path=image_path,
            image_url=image_url,
            reasoning_effort=reasoning_effort,
            raw_schema=raw_schema,
        )

    def invoke_embedding(
        self,
        *,
        input_texts: List[str],
        model: Optional[str] = None,
    ) -> Tuple[List[List[float]], Dict[str, Any], Any]:
        """
        Generate embeddings for the given input texts.

        Args:
            input_texts: List of text inputs
            model: Model identifier (uses default if not specified)

        Returns:
            Tuple of (embeddings, usage_metadata, raw_response)
        """
        model = model or self.default_model
        return self.provider.invoke_embedding(
            input_texts=input_texts,
            model=model,
        )

    def get_default_model(self) -> str:
        return self.default_model

    # ==========================================================================================================
    # BATCH PROCESSING METHODS
    # ==========================================================================================================

    def get_batch_status(self, batch_id: str) -> Dict[str, Any]:
        """
        Get the current status of a batch job.

        Args:
            batch_id: ID of the batch

        Returns:
            Batch object with current status
        """
        return self.provider.get_batch_status(batch_id)

    def wait_for_batch(
        self,
        batch_id: str,
        poll_interval: float = 30.0,
        timeout: Optional[float] = None,
        callback: Optional[Callable[[Dict[str, Any]], None]] = None,
    ) -> Dict[str, Any]:
        """
        Wait for a batch job to complete, polling periodically.

        Args:
            batch_id: ID of the batch
            poll_interval: Seconds between status checks (default: 30)
            timeout: Maximum seconds to wait (default: None = no timeout)
            callback: Optional function called with status on each poll

        Returns:
            Final batch status

        Raises:
            TimeoutError: If timeout is reached before completion
            RuntimeError: If batch fails or is cancelled
        """
        return self.provider.wait_for_batch(batch_id, poll_interval, timeout, callback)

    def get_batch_results(
        self,
        batch_id: str,
        output_path: Optional[str] = None,
    ) -> List[Dict[str, Any]]:
        """
        Retrieve results from a completed batch.

        Args:
            batch_id: ID of the batch
            output_path: Optional path to save the raw JSONL output

        Returns:
            List of result dictionaries, each containing:
                - custom_id: The original request ID
                - output_text: The generated text (flattened)
                - error: Error data (if failed)
        """
        return self.provider.get_batch_results(batch_id, output_path)

    def get_partial_batch_results(self, batch_id: str) -> List[Dict[str, Any]]:
        """
        Retrieve available results from a batch regardless of completion status.

        Unlike get_batch_results, this does not raise on non-completed batches,
        making it suitable for expired batches with partial output.

        Args:
            batch_id: ID of the batch

        Returns:
            List of flattened result dicts (custom_id, output_text, error).
        """
        return self.provider.get_partial_batch_results(batch_id)

    def get_batch_errors(self, batch_id: str) -> List[Dict[str, Any]]:
        """
        Retrieve errors from a batch (if any).

        Args:
            batch_id: ID of the batch

        Returns:
            List of error dictionaries
        """
        return self.provider.get_batch_errors(batch_id)

    def cancel_batch(self, batch_id: str) -> Dict[str, Any]:
        """
        Cancel a running batch job.

        Args:
            batch_id: ID of the batch to cancel

        Returns:
            Updated batch status
        """
        return self.provider.cancel_batch(batch_id)

    def list_batches(self, limit: int = 20, after: Optional[str] = None) -> List[Dict[str, Any]]:
        """
        List all batches.

        Args:
            limit: Maximum number of batches to return
            after: Cursor for pagination

        Returns:
            List of batch objects
        """
        return self.provider.list_batches(limit, after)

    def submit_batch(
        self,
        requests: List[Dict[str, Any]],
        metadata: Optional[Dict[str, str]] = None,
        wait: bool = False,
        poll_interval: float = 30.0,
        callback: Optional[Callable[[Dict[str, Any]], None]] = None,
    ) -> Tuple[str, Optional[List[Dict[str, Any]]]]:
        """
        High-level method to submit a batch job (create file, upload, and start batch).

        Args:
            requests: List of request configurations. Each dict should have:
                - custom_id: Unique identifier
                - instructions: System instructions
                - input_texts: List of input texts
                - model: Model to use (optional, uses default)
                - image_url: Image URL (optional)
                - image_path: Local image path (optional)
                - reasoning_effort: Reasoning effort level (optional)
                - raw_schema: Raw JSON schema dict (optional)
            metadata: Optional metadata for the batch
            wait: If True, wait for completion and return results
            poll_interval: Seconds between status checks when waiting
            callback: Optional callback for status updates when waiting

        Returns:
            Tuple of (batch_id, results). Results is None if wait=False.
        """
        return self.provider.submit_batch(requests, metadata, wait, poll_interval, callback)
