/headroom/proxy/handlers/batch.py

"""Batch handler mixin for HeadroomProxy.

Contains all batch API handlers for Google and OpenAI batch operations.
"""

from __future__ import annotations

import json
import logging
import time
from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
    from fastapi import Request
    from fastapi.responses import Response

from headroom.proxy.auth_mode import classify_client
from headroom.proxy.helpers import COMPRESSION_TIMEOUT_SECONDS, extract_tags
from headroom.proxy.outcome import RequestOutcome

logger = logging.getLogger("headroom.proxy")


class BatchHandlerMixin:
    """Mixin providing batch API handler methods for HeadroomProxy."""

    async def handle_google_batch_create(
        self,
        request: Request,
        model: str,
    ) -> Response:
        """Handle Google POST /v1beta/models/{model}:batchGenerateContent endpoint.

        Google batch format:
        {
            "batch": {
                "display_name": "my-batch",
                "input_config": {
                    "requests": {
                        "requests": [
                            {
                                "request": {"contents": [{"parts": [{"text": "..."}]}]},
                                "metadata": {"key": "request-1"}
                            }
                        ]
                    }
                }
            }
        }

        This method applies compression to each request's contents before forwarding.
        """
        from fastapi.responses import JSONResponse, Response

        from headroom.ccr import CCRToolInjector
        from headroom.proxy.helpers import MAX_REQUEST_BODY_SIZE, _read_request_json
        from headroom.utils import extract_user_query

        start_time = time.time()
        request_id = await self._next_request_id()

        # Check request body size
        content_length = request.headers.get("content-length")
        if content_length and int(content_length) > MAX_REQUEST_BODY_SIZE:
            return JSONResponse(
                status_code=413,
                content={
                    "error": {
                        "code": 413,
                        "message": f"Request body too large. Maximum size is {MAX_REQUEST_BODY_SIZE // (1024 * 1024)}MB",
                        "status": "INVALID_ARGUMENT",
                    }
                },
            )

        # Parse request
        try:
            body = await _read_request_json(request)
        except (json.JSONDecodeError, ValueError) as e:
            return JSONResponse(
                status_code=400,
                content={
                    "error": {
                        "code": 400,
                        "message": f"Invalid request body: {e!s}",
                        "status": "INVALID_ARGUMENT",
                    }
                },
            )

        # Extract batch config
        batch_config = body.get("batch", {})
        input_config = batch_config.get("input_config", {})
        requests_wrapper = input_config.get("requests", {})
        requests_list = requests_wrapper.get("requests", [])

        if not requests_list:
            # No inline requests - might be using file input, pass through
            logger.debug(f"[{request_id}] Google batch: No inline requests, passing through")
            return await self._google_batch_passthrough(request, model, body)

        # Extract headers
        headers = dict(request.headers.items())
        headers.pop("host", None)
        headers.pop("content-length", None)
        client = classify_client(headers)
        tags = extract_tags(headers)
        # PR-A5 (P5-49): strip internal x-headroom-* before forwarding upstream.
        from headroom.proxy.helpers import _strip_internal_headers, log_outbound_headers

        _pre_strip_count_gb = sum(1 for k in headers if k.lower().startswith("x-headroom-"))
        headers = _strip_internal_headers(headers)
        log_outbound_headers(
            forwarder="google_batch",
            stripped_count=_pre_strip_count_gb,
            request_id=request_id,
        )

        # Track compression stats
        total_original_tokens = 0
        total_optimized_tokens = 0
        total_tokens_saved = 0
        compressed_requests = []
        pipeline_timing: dict[str, float] = {}

        # Apply compression to each request in the batch
        for idx, batch_req in enumerate(requests_list):
            req_content = batch_req.get("request", {})
            metadata = batch_req.get("metadata", {})
            contents = req_content.get("contents", [])

            if not contents or not self.config.optimize:
                # No contents or optimization disabled - pass through unchanged
                compressed_requests.append(batch_req)
                continue

            # Convert Google format to messages for compression
            system_instruction = req_content.get("systemInstruction")
            messages, preserved_indices = self._gemini_contents_to_messages(
                contents, system_instruction
            )

            # Store original content entries that have non-text parts before compression
            preserved_contents = {idx: contents[idx] for idx in preserved_indices}

            # Early exit if ALL content has non-text parts (nothing to compress)
            if len(preserved_indices) == len(contents):
                # All content has non-text parts, skip compression
                compressed_requests.append(batch_req)
                continue

            # Apply optimization
            original_tokens = 0  # Set before try so error handler can use it
            optimized_tokens = 0
            try:
                # Look up model context limit, fall back to 128K
                context_limit = (
                    self.openai_provider.get_context_limit(model)
                    if hasattr(self, "openai_provider")
                    else 128000
                )

                # Use OpenAI pipeline (similar message format after conversion)
                # Offload off the event loop (#1701): inline apply() blocks
                # every other request; timeouts fall to the except below.
                result = await self._run_compression_in_executor(
                    lambda messages=messages, model=model, context_limit=context_limit: (
                        self.openai_pipeline.apply(
                            messages=messages,
                            model=model,
                            model_limit=context_limit,
                            context=extract_user_query(messages),
                        )
                    ),
                    timeout=COMPRESSION_TIMEOUT_SECONDS,
                )

                optimized_messages = result.messages
                for k, v in result.timing.items():
                    pipeline_timing[k] = pipeline_timing.get(k, 0.0) + v
                # Use pipeline's token counts for consistency with pipeline logs
                original_tokens = result.tokens_before
                optimized_tokens = result.tokens_after
                # Guard: if "optimization" inflated tokens, revert to originals
                if optimized_tokens > original_tokens:
                    logger.warning(
                        f"[{request_id}] Batch item optimization inflated tokens "
                        f"({original_tokens} -> {optimized_tokens}), reverting"
                    )
                    optimized_messages = messages
                    optimized_tokens = original_tokens

                total_original_tokens += original_tokens
                total_optimized_tokens += optimized_tokens
                tokens_saved = original_tokens - optimized_tokens
                total_tokens_saved += tokens_saved

                # CCR Tool Injection: Inject retrieval tool if compression occurred
                tools = req_content.get("tools")
                # Extract existing function declarations if present
                existing_funcs = None
                if tools:
                    for tool in tools:
                        if "functionDeclarations" in tool:
                            existing_funcs = tool["functionDeclarations"]
                            break

                if self.config.ccr_inject_tool and tokens_saved > 0:
                    injector = CCRToolInjector(
                        provider="google",
                        inject_tool=True,
                        inject_system_instructions=self.config.ccr_inject_system_instructions,
                    )
                    optimized_messages, injected_funcs, was_injected = injector.process_request(
                        optimized_messages, existing_funcs
                    )
                    if was_injected:
                        logger.debug(
                            f"[{request_id}] CCR: Injected retrieval tool for Google batch request {idx}"
                        )
                        existing_funcs = injected_funcs

                # Convert back to Google contents format
                optimized_contents, optimized_sys_inst = self._messages_to_gemini_contents(
                    optimized_messages
                )

                # Restore preserved content entries that had non-text parts
                for orig_idx, original_content in preserved_contents.items():
                    if orig_idx < len(optimized_contents):
                        optimized_contents[orig_idx] = original_content

                # Create compressed batch request
                compressed_req_content = {**req_content, "contents": optimized_contents}
                if optimized_sys_inst:
                    compressed_req_content["systemInstruction"] = optimized_sys_inst
                if existing_funcs is not None:
                    compressed_req_content["tools"] = [{"functionDeclarations": existing_funcs}]

                compressed_req = {
                    "request": compressed_req_content,
                    "metadata": metadata,
                }

                compressed_requests.append(compressed_req)

                if tokens_saved > 0:
                    logger.debug(
                        f"[{request_id}] Google batch request {idx}: "
                        f"{original_tokens:,} -> {optimized_tokens:,} tokens "
                        f"(saved {tokens_saved:,})"
                    )

            except Exception as e:
                logger.warning(
                    f"[{request_id}] Optimization failed for Google batch request {idx}: {e}"
                )
                # Pass through unchanged on failure — count original as optimized
                compressed_requests.append(batch_req)
                total_optimized_tokens += original_tokens  # 0 if pipeline never ran

        # Update body with compressed requests
        body["batch"]["input_config"]["requests"]["requests"] = compressed_requests

        optimization_latency = (time.time() - start_time) * 1000

        # Forward request to Google
        url = f"{self.GEMINI_API_URL}/v1beta/models/{model}:batchGenerateContent"

        # Add API key to URL if present in headers
        api_key = headers.pop("x-goog-api-key", None)
        if api_key:
            url = f"{url}?key={api_key}"

        try:
            response = await self._retry_request("POST", url, headers, body)

            # Google batch create — funnel records via the canonical
            # path; cache fields stay 0 (Google batches don't expose
            # cache stats in the same shape).
            await self._record_request_outcome(
                RequestOutcome(
                    request_id=request_id,
                    provider="google",
                    model=f"batch:{model}",
                    original_tokens=total_original_tokens,
                    optimized_tokens=total_optimized_tokens,
                    output_tokens=0,
                    tokens_saved=total_tokens_saved,
                    attempted_input_tokens=total_optimized_tokens + total_tokens_saved,
                    total_latency_ms=optimization_latency,
                    overhead_ms=optimization_latency,
                    pipeline_timing=pipeline_timing,
                    tags=tags,
                    client=client,
                )
            )

            # Log compression stats
            if total_tokens_saved > 0:
                savings_percent = (
                    (total_tokens_saved / total_original_tokens * 100)
                    if total_original_tokens > 0
                    else 0
                )
                logger.info(
                    f"[{request_id}] Google batch compression: "
                    f"{total_original_tokens:,} -> {total_optimized_tokens:,} tokens "
                    f"({savings_percent:.1f}% saved across {len(requests_list)} requests)"
                )

            # Store batch context for CCR result processing
            if response.status_code == 200 and self.config.ccr_inject_tool:
                try:
                    response_data = response.json()
                    batch_name = response_data.get("name")
                    if batch_name:
                        await self._store_google_batch_context(
                            batch_name,
                            requests_list,
                            model,
                            api_key,
                        )
                except Exception as e:
                    logger.warning(f"[{request_id}] Failed to store Google batch context: {e}")

            # Remove compression headers
            response_headers = dict(response.headers)
            response_headers.pop("content-encoding", None)
            response_headers.pop("content-length", None)

            return Response(
                content=response.content,
                status_code=response.status_code,
                headers=response_headers,
            )

        except Exception as e:
            logger.error(f"[{request_id}] Google batch request failed: {e}")
            return JSONResponse(
                status_code=500,
                content={
                    "error": {
                        "code": 500,
                        "message": f"Failed to forward batch request: {e!s}",
                        "status": "INTERNAL",
                    }
                },
            )

    async def _google_batch_passthrough(
        self,
        request: Request,
        model: str,
        body: dict | None = None,
    ) -> Response:
        """Pass through Google batch request without modification."""
        from fastapi.responses import Response

        start_time = time.time()

        headers = dict(request.headers.items())
        headers.pop("host", None)
        headers.pop("content-length", None)
        client = classify_client(headers)
        tags = extract_tags(headers)
        # PR-A5 (P5-49): strip internal x-headroom-* before forwarding upstream.
        from headroom.proxy.helpers import _strip_internal_headers, log_outbound_headers

        _pre_strip_count_gpt = sum(1 for k in headers if k.lower().startswith("x-headroom-"))
        headers = _strip_internal_headers(headers)
        log_outbound_headers(
            forwarder="google_batch_passthrough",
            stripped_count=_pre_strip_count_gpt,
            request_id=None,
        )

        url = f"{self.GEMINI_API_URL}/v1beta/models/{model}:batchGenerateContent"

        # Add API key to URL if present in headers
        api_key = headers.pop("x-goog-api-key", None)
        if api_key:
            url = f"{url}?key={api_key}"

        # Byte-faithful body bytes (PR-A3, fixes P0-2). When ``body`` is
        # None we forward the original bytes verbatim; otherwise the dict
        # has been synthesized by Headroom and is canonically serialized.
        from headroom.proxy.body_forwarding import (
            get_python_forwarder_mode,
            prepare_outbound_body_bytes,
            serialize_body_canonical,
        )
        from headroom.proxy.helpers import log_outbound_request

        if body is None:
            from starlette.requests import ClientDisconnect

            try:
                body_content = await request.body()
            except ClientDisconnect:
                logger.debug("Client disconnected during body read for google batch passthrough")
                return Response(status_code=204)
            outbound_source = "passthrough"
            body_mutated = False
        else:
            body_content = serialize_body_canonical(body)
            outbound_source = "canonical"
            body_mutated = True
        log_outbound_request(
            forwarder="google_batch_passthrough",
            method="POST",
            path=url,
            body_bytes_count=len(body_content),
            body_mutated=body_mutated,
            mutation_reasons=["google_batch_resynthesized"] if body_mutated else [],
            request_id=None,
            source=outbound_source,
        )
        # ``prepare_outbound_body_bytes`` is consulted only for the legacy
        # operator opt-in path so we honor the env-var override here too.
        if get_python_forwarder_mode() == "legacy_json_kwarg" and body is not None:
            outbound_bytes, _ = prepare_outbound_body_bytes(
                body=body,
                original_body_bytes=None,
                body_mutated=True,
            )
            body_content = outbound_bytes

        response = await self.http_client.post(  # type: ignore[union-attr]
            url,
            headers=headers,
            content=body_content,
        )

        # Google batch (Files API forward) — no compression, just
        # upstream forward. Funnel records via zero defaults so the
        # request shows up in dashboards even with no token activity.
        latency_ms = (time.time() - start_time) * 1000
        request_id_files = await self._next_request_id()
        await self._record_request_outcome(
            RequestOutcome(
                request_id=request_id_files,
                provider="google",
                model=f"passthrough:batch:{model}",
                original_tokens=0,
                optimized_tokens=0,
                output_tokens=0,
                tokens_saved=0,
                attempted_input_tokens=0,
                total_latency_ms=latency_ms,
                tags=tags,
                client=client,
            )
        )

        response_headers = dict(response.headers)
        response_headers.pop("content-encoding", None)
        response_headers.pop("content-length", None)

        return Response(
            content=response.content,
            status_code=response.status_code,
            headers=response_headers,
        )

    async def handle_google_batch_passthrough(
        self,
        request: Request,
        batch_name: str | None = None,
    ) -> Response:
        """Handle Google batch passthrough endpoints.

        Used for:
        - GET /v1beta/batches/{batch_name} - Get batch status
        - POST /v1beta/batches/{batch_name}:cancel - Cancel batch
        - DELETE /v1beta/batches/{batch_name} - Delete batch
        """
        from fastapi.responses import Response

        start_time = time.time()
        path = request.url.path
        url = f"{self.GEMINI_API_URL}{path}"

        # Preserve query string parameters
        if request.url.query:
            url = f"{url}?{request.url.query}"

        headers = dict(request.headers.items())
        headers.pop("host", None)
        client = classify_client(headers)
        tags = extract_tags(headers)
        # PR-A5 (P5-49): strip internal x-headroom-* before forwarding upstream.
        from headroom.proxy.helpers import _strip_internal_headers, log_outbound_headers

        _pre_strip_count_gp = sum(1 for k in headers if k.lower().startswith("x-headroom-"))
        headers = _strip_internal_headers(headers)
        log_outbound_headers(
            forwarder="gemini_passthrough",
            stripped_count=_pre_strip_count_gp,
            request_id=None,
        )

        # Handle API key
        api_key = headers.pop("x-goog-api-key", None)
        if api_key:
            if "?" in url:
                url = f"{url}&key={api_key}"
            else:
                url = f"{url}?key={api_key}"

        from starlette.requests import ClientDisconnect

        try:
            body = await request.body()
        except ClientDisconnect:
            logger.debug("Client disconnected during body read for gemini passthrough")
            return Response(status_code=204)

        response = await self.http_client.request(  # type: ignore[union-attr]
            method=request.method,
            url=url,
            headers=headers,
            content=body,
        )

        # Google batch passthrough — list/get/cancel forwarded with
        # no compression work. Funnel records the request so it's
        # visible in dashboards.
        latency_ms = (time.time() - start_time) * 1000
        request_id = await self._next_request_id()
        await self._record_request_outcome(
            RequestOutcome(
                request_id=request_id,
                provider="google",
                model="passthrough:batches",
                original_tokens=0,
                optimized_tokens=0,
                output_tokens=0,
                tokens_saved=0,
                attempted_input_tokens=0,
                total_latency_ms=latency_ms,
                tags=tags,
                client=client,
            )
        )

        response_headers = dict(response.headers)
        response_headers.pop("content-encoding", None)
        response_headers.pop("content-length", None)

        return Response(
            content=response.content,
            status_code=response.status_code,
            headers=response_headers,
        )

    async def _store_google_batch_context(
        self,
        batch_name: str,
        requests_list: list[dict[str, Any]],
        model: str,
        api_key: str | None,
    ) -> None:
        """Store Google batch context for CCR result processing.

        Args:
            batch_name: The batch name from the API response.
            requests_list: The original batch requests.
            model: The model used for the batch.
            api_key: The API key for continuation calls.
        """
        from headroom.ccr import BatchContext, BatchRequestContext, get_batch_context_store

        store = get_batch_context_store()
        context = BatchContext(
            batch_id=batch_name,
            provider="google",
            api_key=api_key,
            api_base_url=self.GEMINI_API_URL,
        )

        for batch_req in requests_list:
            metadata = batch_req.get("metadata", {})
            custom_id = metadata.get("key", "")
            req_content = batch_req.get("request", {})
            contents = req_content.get("contents", [])
            system_instruction = req_content.get("systemInstruction")

            # Convert contents to messages format for CCR handler
            messages, _ = self._gemini_contents_to_messages(contents, system_instruction)

            # Extract system instruction text if present
            sys_text = None
            if system_instruction:
                parts = system_instruction.get("parts", [])
                if parts and isinstance(parts[0], dict):
                    sys_text = parts[0].get("text")

            context.add_request(
                BatchRequestContext(
                    custom_id=custom_id,
                    messages=messages,
                    tools=req_content.get("tools"),
                    model=model,
                    system_instruction=sys_text,
                )
            )

        await store.store(context)
        logger.debug(
            f"Stored Google batch context for {batch_name} with {len(requests_list)} requests"
        )

    async def handle_google_batch_results(
        self,
        request: Request,
        batch_name: str,
    ) -> Response:
        """Handle Google batch results with CCR post-processing.

        Google batch results endpoint returns the batch operation status.
        When status is SUCCEEDED, results are embedded in the response.
        This handler processes CCR tool calls in those results.
        """
        from fastapi.responses import JSONResponse, Response

        from headroom.ccr import BatchResultProcessor, get_batch_context_store

        start_time = time.time()

        # Forward request to get batch status/results
        url = f"{self.GEMINI_API_URL}/v1beta/{batch_name}"

        if request.url.query:
            url = f"{url}?{request.url.query}"

        headers = dict(request.headers.items())
        headers.pop("host", None)
        client = classify_client(headers)
        tags = extract_tags(headers)
        # PR-A5 (P5-49): strip internal x-headroom-* before forwarding upstream.
        from headroom.proxy.helpers import _strip_internal_headers, log_outbound_headers

        _pre_strip_count_gbr = sum(1 for k in headers if k.lower().startswith("x-headroom-"))
        headers = _strip_internal_headers(headers)
        log_outbound_headers(
            forwarder="google_batch_results",
            stripped_count=_pre_strip_count_gbr,
            request_id=None,
        )

        # Handle API key
        api_key = headers.pop("x-goog-api-key", None)
        if api_key:
            if "?" in url:
                url = f"{url}&key={api_key}"
            else:
                url = f"{url}?key={api_key}"

        response = await self.http_client.get(url, headers=headers)  # type: ignore[union-attr]

        if response.status_code != 200:
            # Error - pass through
            response_headers = dict(response.headers)
            response_headers.pop("content-encoding", None)
            response_headers.pop("content-length", None)
            return Response(
                content=response.content,
                status_code=response.status_code,
                headers=response_headers,
            )

        # Parse response
        try:
            response_data = response.json()
        except json.JSONDecodeError:
            # Not JSON - pass through
            response_headers = dict(response.headers)
            response_headers.pop("content-encoding", None)
            response_headers.pop("content-length", None)
            return Response(
                content=response.content,
                status_code=response.status_code,
                headers=response_headers,
            )

        # Check if batch has results (state must be SUCCEEDED)
        metadata = response_data.get("metadata", {})
        state = metadata.get("state")

        if state != "SUCCEEDED":
            # Batch not complete - pass through
            response_headers = dict(response.headers)
            response_headers.pop("content-encoding", None)
            response_headers.pop("content-length", None)
            return Response(
                content=response.content,
                status_code=response.status_code,
                headers=response_headers,
            )

        # Extract results from response
        # Google embeds results in the batch response
        results = response_data.get("response", {}).get("responses", [])

        if not results:
            # No results to process
            response_headers = dict(response.headers)
            response_headers.pop("content-encoding", None)
            response_headers.pop("content-length", None)
            return Response(
                content=response.content,
                status_code=response.status_code,
                headers=response_headers,
            )

        # Check if we have context and CCR processing is enabled
        store = get_batch_context_store()
        batch_context = await store.get(batch_name)

        if batch_context is None or not self.config.ccr_inject_tool:
            # No context or CCR disabled - pass through
            response_headers = dict(response.headers)
            response_headers.pop("content-encoding", None)
            response_headers.pop("content-length", None)
            return Response(
                content=response.content,
                status_code=response.status_code,
                headers=response_headers,
            )

        # Process results with CCR handler
        processor = BatchResultProcessor(self.http_client)  # type: ignore[arg-type]
        processed = await processor.process_results(batch_name, results, "google")

        # Update response with processed results
        processed_results = [p.result for p in processed]
        response_data["response"]["responses"] = processed_results

        for p in processed:
            if p.was_processed:
                logger.info(
                    f"CCR: Processed Google batch result {p.custom_id} "
                    f"({p.continuation_rounds} continuation rounds)"
                )

        # Google batch results with CCR processing — funnel records
        # via zero defaults.
        latency_ms = (time.time() - start_time) * 1000
        request_id = await self._next_request_id()
        await self._record_request_outcome(
            RequestOutcome(
                request_id=request_id,
                provider="google",
                model="batch:ccr-processed",
                original_tokens=0,
                optimized_tokens=0,
                output_tokens=0,
                tokens_saved=0,
                attempted_input_tokens=0,
                total_latency_ms=latency_ms,
                tags=tags,
                client=client,
            )
        )

        return JSONResponse(content=response_data, status_code=200)

    async def handle_batch_create(self, request: Request) -> Response:
        """Handle POST /v1/batches - Create a batch with compression.

        Flow:
        1. Parse request to get input_file_id
        2. Download the JSONL file content from OpenAI
        3. Parse each line and compress the messages
        4. Create a new compressed JSONL file
        5. Upload compressed file to OpenAI
        6. Create batch with the new compressed file_id
        7. Return batch object with compression stats in metadata
        """
        from fastapi.responses import JSONResponse, Response

        from headroom.proxy.helpers import _read_request_json

        start_time = time.time()
        request_id = await self._next_request_id()

        # Parse request
        try:
            body = await _read_request_json(request)
        except (json.JSONDecodeError, ValueError) as e:
            return JSONResponse(
                status_code=400,
                content={
                    "error": {
                        "message": f"Invalid request body: {e!s}",
                        "type": "invalid_request_error",
                        "code": "invalid_json",
                    }
                },
            )

        input_file_id = body.get("input_file_id")
        endpoint = body.get("endpoint")
        completion_window = body.get("completion_window", "24h")
        metadata = body.get("metadata", {})

        if not input_file_id:
            return JSONResponse(
                status_code=400,
                content={
                    "error": {
                        "message": "input_file_id is required",
                        "type": "invalid_request_error",
                        "code": "missing_parameter",
                    }
                },
            )

        if not endpoint:
            return JSONResponse(
                status_code=400,
                content={
                    "error": {
                        "message": "endpoint is required",
                        "type": "invalid_request_error",
                        "code": "missing_parameter",
                    }
                },
            )

        # Only compress chat completions endpoint
        if endpoint != "/v1/chat/completions":
            # Pass through for other endpoints
            return await self._batch_passthrough(request, body)

        headers = dict(request.headers.items())
        headers.pop("host", None)
        headers.pop("content-length", None)
        client = classify_client(headers)
        tags = extract_tags(headers)
        # PR-A5 (P5-49): strip internal x-headroom-* before forwarding upstream.
        from headroom.proxy.helpers import _strip_internal_headers, log_outbound_headers

        _pre_strip_count_oacc = sum(1 for k in headers if k.lower().startswith("x-headroom-"))
        headers = _strip_internal_headers(headers)
        log_outbound_headers(
            forwarder="openai_batch_chat_completions",
            stripped_count=_pre_strip_count_oacc,
            request_id=request_id,
        )

        try:
            # Step 1: Download the input file from OpenAI
            logger.info(f"[{request_id}] Batch: Downloading input file {input_file_id}")
            file_content = await self._download_openai_file(input_file_id, headers)

            if file_content is None:
                return JSONResponse(
                    status_code=404,
                    content={
                        "error": {
                            "message": f"Failed to download file {input_file_id}",
                            "type": "invalid_request_error",
                            "code": "file_not_found",
                        }
                    },
                )

            # Step 2: Parse and compress each line
            logger.info(f"[{request_id}] Batch: Compressing JSONL content")
            compressed_lines, stats = await self._compress_batch_jsonl(file_content, request_id)

            if stats["total_requests"] == 0:
                return JSONResponse(
                    status_code=400,
                    content={
                        "error": {
                            "message": "No valid requests found in input file",
                            "type": "invalid_request_error",
                            "code": "empty_file",
                        }
                    },
                )

            # Step 3: Create compressed JSONL content
            compressed_content = "\n".join(compressed_lines)

            # Step 4: Upload compressed file to OpenAI
            logger.info(f"[{request_id}] Batch: Uploading compressed file")
            new_file_id = await self._upload_openai_file(
                compressed_content, f"compressed_{input_file_id}.jsonl", headers
            )

            if new_file_id is None:
                return JSONResponse(
                    status_code=500,
                    content={
                        "error": {
                            "message": "Failed to upload compressed file",
                            "type": "server_error",
                            "code": "upload_failed",
                        }
                    },
                )

            # Step 5: Create batch with compressed file
            logger.info(f"[{request_id}] Batch: Creating batch with compressed file {new_file_id}")

            # Add compression stats to metadata
            compression_metadata = {
                **metadata,
                "headroom_compressed": "true",
                "headroom_original_file_id": input_file_id,
                "headroom_total_requests": str(stats["total_requests"]),
                "headroom_tokens_saved": str(stats["total_tokens_saved"]),
                "headroom_original_tokens": str(stats["total_original_tokens"]),
                "headroom_compressed_tokens": str(stats["total_compressed_tokens"]),
                "headroom_savings_percent": f"{stats['savings_percent']:.1f}",
            }

            batch_body = {
                "input_file_id": new_file_id,
                "endpoint": endpoint,
                "completion_window": completion_window,
                "metadata": compression_metadata,
            }

            url = f"{self.OPENAI_API_URL}/v1/batches"

            # Byte-faithful re-serialization (PR-A3, fixes P0-2).
            # batch_body is synthesized by Headroom (compressed file_id +
            # metadata), so it is treated as mutated and goes through the
            # canonical serializer.
            from headroom.proxy.body_forwarding import prepare_outbound_body_bytes
            from headroom.proxy.helpers import log_outbound_request

            outbound_bytes, outbound_source = prepare_outbound_body_bytes(
                body=batch_body,
                original_body_bytes=None,
                body_mutated=True,
            )
            outbound_headers = {**headers, "content-type": "application/json"}
            log_outbound_request(
                forwarder="batch",
                method="POST",
                path=url,
                body_bytes_count=len(outbound_bytes),
                body_mutated=True,
                mutation_reasons=["batch_compressed_file_substitution"],
                request_id=request_id,
                source=outbound_source,
            )
            response = await self.http_client.post(  # type: ignore[union-attr]
                url, content=outbound_bytes, headers=outbound_headers
            )

            total_latency = (time.time() - start_time) * 1000

            # Log compression stats
            logger.info(
                f"[{request_id}] Batch created: {stats['total_requests']} requests, "
                f"{stats['total_original_tokens']:,} -> {stats['total_compressed_tokens']:,} tokens "
                f"(saved {stats['total_tokens_saved']:,} tokens, {stats['savings_percent']:.1f}%) "
                f"in {total_latency:.0f}ms"
            )

            # OpenAI batch create — funnel records via the canonical
            # path; `model="batch"` matches the synthetic naming used
            # by the Anthropic batch handlers.
            await self._record_request_outcome(
                RequestOutcome(
                    request_id=request_id,
                    provider="openai",
                    model="batch",
                    original_tokens=stats["total_original_tokens"],
                    optimized_tokens=stats["total_compressed_tokens"],
                    output_tokens=0,
                    tokens_saved=stats["total_tokens_saved"],
                    attempted_input_tokens=stats["total_compressed_tokens"]
                    + stats["total_tokens_saved"],
                    total_latency_ms=total_latency,
                    tags=tags,
                    client=client,
                )
            )

            # Return response with compression info in headers
            response_headers = dict(response.headers)
            response_headers.pop("content-encoding", None)
            response_headers.pop("content-length", None)
            response_headers["x-headroom-tokens-saved"] = str(stats["total_tokens_saved"])
            response_headers["x-headroom-savings-percent"] = f"{stats['savings_percent']:.1f}"

            return Response(
                content=response.content,
                status_code=response.status_code,
                headers=response_headers,
            )

        except Exception as e:
            logger.error(f"[{request_id}] Batch creation failed: {type(e).__name__}: {e}")
            await self.metrics.record_failed(provider="batch")
            return JSONResponse(
                status_code=500,
                content={
                    "error": {
                        "message": "An error occurred while processing the batch request",
                        "type": "server_error",
                        "code": "batch_processing_error",
                    }
                },
            )

    async def _download_openai_file(self, file_id: str, headers: dict) -> str | None:
        """Download file content from OpenAI."""
        url = f"{self.OPENAI_API_URL}/v1/files/{file_id}/content"
        try:
            response = await self.http_client.get(url, headers=headers)  # type: ignore[union-attr]
            if response.status_code == 200:
                return str(response.text)
            logger.error(f"Failed to download file {file_id}: {response.status_code}")
            return None
        except Exception as e:
            logger.error(f"Error downloading file {file_id}: {e}")
            return None

    async def _upload_openai_file(self, content: str, filename: str, headers: dict) -> str | None:
        """Upload a file to OpenAI for batch processing."""
        url = f"{self.OPENAI_API_URL}/v1/files"

        # Prepare multipart form data
        # We need to use httpx's files parameter for multipart upload
        files = {
            "file": (filename, content.encode("utf-8"), "application/jsonl"),
        }
        data = {
            "purpose": "batch",
        }

        # Remove content-type from headers (httpx will set it for multipart)
        upload_headers = {k: v for k, v in headers.items() if k.lower() != "content-type"}

        try:
            response = await self.http_client.post(  # type: ignore[union-attr]
                url, files=files, data=data, headers=upload_headers
            )
            if response.status_code == 200:
                result = response.json()
                file_id: str | None = result.get("id")
                return file_id
            logger.error(f"Failed to upload file: {response.status_code} - {response.text}")
            return None
        except Exception as e:
            logger.error(f"Error uploading file: {e}")
            return None

    async def _compress_batch_jsonl(self, content: str, request_id: str) -> tuple[list[str], dict]:
        """Compress messages in each line of a batch JSONL file.

        Returns:
            Tuple of (compressed_lines, stats_dict)
        """
        from headroom.ccr import CCRToolInjector
        from headroom.tokenizers import get_tokenizer
        from headroom.utils import extract_user_query

        lines = content.strip().split("\n")
        compressed_lines = []
        total_original_tokens = 0
        total_compressed_tokens = 0
        total_requests = 0
        errors = 0

        tokenizer = get_tokenizer("gpt-4")  # Use gpt-4 tokenizer for batch

        for i, line in enumerate(lines):
            if not line.strip():
                continue

            try:
                request_obj = json.loads(line)
                body = request_obj.get("body", {})
                messages = body.get("messages", [])
                model = body.get("model", "gpt-4")

                if not messages:
                    # No messages to compress, pass through
                    compressed_lines.append(line)
                    total_requests += 1
                    continue

                # Compress messages using the OpenAI pipeline
                if self.config.optimize:
                    try:
                        context_limit = self.openai_provider.get_context_limit(model)
                        # Offload off the event loop (#1701); timeouts fall to
                        # the except below and pass the line through.
                        result = await self._run_compression_in_executor(
                            lambda messages=messages, model=model, context_limit=context_limit: (
                                self.openai_pipeline.apply(
                                    messages=messages,
                                    model=model,
                                    model_limit=context_limit,
                                    context=extract_user_query(messages),
                                )
                            ),
                            timeout=COMPRESSION_TIMEOUT_SECONDS,
                        )
                        compressed_messages = result.messages
                        # Use pipeline's token counts for consistency with pipeline logs
                        original_tokens = result.tokens_before
                        compressed_tokens = result.tokens_after
                    except Exception as e:
                        logger.warning(f"[{request_id}] Compression failed for line {i}: {e}")
                        compressed_messages = messages
                        original_tokens = tokenizer.count_messages(messages)
                        compressed_tokens = original_tokens
                else:
                    compressed_messages = messages
                    original_tokens = tokenizer.count_messages(messages)
                    compressed_tokens = original_tokens

                total_original_tokens += original_tokens
                total_compressed_tokens += compressed_tokens
                tokens_saved = original_tokens - compressed_tokens

                # CCR Tool Injection: Inject retrieval tool if compression occurred
                tools = body.get("tools")
                if self.config.ccr_inject_tool and tokens_saved > 0:
                    injector = CCRToolInjector(
                        provider="openai",
                        inject_tool=True,
                        inject_system_instructions=self.config.ccr_inject_system_instructions,
                    )
                    compressed_messages, tools, was_injected = injector.process_request(
                        compressed_messages, tools
                    )
                    if was_injected:
                        logger.debug(
                            f"[{request_id}] CCR: Injected retrieval tool for batch line {i}"
                        )

                # Update body with compressed messages
                body["messages"] = compressed_messages
                if tools is not None:
                    body["tools"] = tools
                request_obj["body"] = body

                compressed_lines.append(json.dumps(request_obj))
                total_requests += 1

            except json.JSONDecodeError as e:
                logger.warning(f"[{request_id}] Invalid JSON on line {i}: {e}")
                errors += 1
                # Keep original line on error
                compressed_lines.append(line)
                total_requests += 1

        total_tokens_saved = total_original_tokens - total_compressed_tokens
        savings_percent = (
            (total_tokens_saved / total_original_tokens * 100) if total_original_tokens > 0 else 0
        )

        stats = {
            "total_requests": total_requests,
            "total_original_tokens": total_original_tokens,
            "total_compressed_tokens": total_compressed_tokens,
            "total_tokens_saved": total_tokens_saved,
            "savings_percent": savings_percent,
            "errors": errors,
        }

        return compressed_lines, stats

    async def _batch_passthrough(self, request: Request, body: dict) -> Response:
        """Pass through batch request to OpenAI without compression.

        Byte-faithful (PR-A3, fixes P0-2). The original request bytes are
        preserved verbatim when no transform mutated the body.
        """
        from fastapi.responses import Response

        from headroom.proxy.body_forwarding import prepare_outbound_body_bytes
        from headroom.proxy.helpers import (
            _read_request_body_bytes,
            _strip_internal_headers,
            log_outbound_headers,
            log_outbound_request,
        )

        headers = dict(request.headers.items())
        headers.pop("host", None)
        headers.pop("content-length", None)
        # PR-A5 (P5-49): strip internal x-headroom-* before forwarding upstream.
        _pre_strip_count_obp = sum(1 for k in headers if k.lower().startswith("x-headroom-"))
        headers = _strip_internal_headers(headers)
        log_outbound_headers(
            forwarder="openai_batch_passthrough",
            stripped_count=_pre_strip_count_obp,
            request_id=None,
        )

        url = f"{self.OPENAI_API_URL}/v1/batches"

        # Best effort: capture the original (decompressed) bytes so the
        # passthrough is truly byte-faithful. If the body was already
        # consumed upstream we fall through to canonical re-serialization.
        try:
            original_body_bytes: bytes | None = await _read_request_body_bytes(request)
        except Exception:
            original_body_bytes = None

        outbound_bytes, outbound_source = prepare_outbound_body_bytes(
            body=body,
            original_body_bytes=original_body_bytes,
            body_mutated=False,
        )
        outbound_headers = {**headers, "content-type": "application/json"}
        log_outbound_request(
            forwarder="batch_passthrough",
            method="POST",
            path=url,
            body_bytes_count=len(outbound_bytes),
            body_mutated=False,
            mutation_reasons=[],
            request_id=None,
            source=outbound_source,
        )
        response = await self.http_client.post(  # type: ignore[union-attr]
            url, content=outbound_bytes, headers=outbound_headers
        )

        response_headers = dict(response.headers)
        response_headers.pop("content-encoding", None)
        response_headers.pop("content-length", None)

        return Response(
            content=response.content,
            status_code=response.status_code,
            headers=response_headers,
        )

    async def handle_batch_list(self, request: Request) -> Response:
        """Handle GET /v1/batches - List batches (passthrough)."""
        return await self.handle_passthrough(request, self.OPENAI_API_URL)

    async def handle_batch_get(self, request: Request, batch_id: str) -> Response:
        """Handle GET /v1/batches/{batch_id} - Get batch (passthrough)."""
        return await self.handle_passthrough(request, self.OPENAI_API_URL)

    async def handle_batch_cancel(self, request: Request, batch_id: str) -> Response:
        """Handle POST /v1/batches/{batch_id}/cancel - Cancel batch (passthrough)."""
        return await self.handle_passthrough(request, self.OPENAI_API_URL)