"""LLM orchestrator for recursive sub-calls within sandbox execution.

Handles recursive LLM sub-calls with depth/call limits and batch processing.
"""

from __future__ import annotations

import json
import sys
from typing import Any, TYPE_CHECKING

from .exceptions import SandboxError, BatchLLMError
from .prompts import SUB_QUERY_PROMPT

if TYPE_CHECKING:
    from .execution_context import ExecutionContext
    from .llm_client import LLMClient


class LlmOrchestrator:
    """Handles recursive LLM sub-calls within sandbox execution.

    Responsibilities:
    - Execute llm_query() sub-calls with depth/call limits
    - Batch processing of multiple queries
    - Context truncation for token limits

    Dependencies are injectable for testing:
    - client: LLMClient implementation for making API calls
    - context: ExecutionContext for tracking state
    """

    MAX_CONTEXT_CHARS: int = 8000
    SUB_CALL_MAX_TOKENS: int = 2048
    BATCH_MAX_PROMPTS: int = 50
    BATCH_MAX_CONCURRENT: int = 5

    def __init__(
        self,
        context: ExecutionContext,
        client: LLMClient,
        subcall_model: str | None = None,
    ) -> None:
        """Initialize the orchestrator.

        Args:
            context: ExecutionContext for tracking depth, calls, and tokens
            client: LLMClient for making API calls
            subcall_model: Model to use for sub-calls (optional)
        """
        self._context = context
        self._client = client
        self._model = subcall_model or self._resolve_subcall_model()

    def _resolve_subcall_model(self) -> str:
        """Resolve the subcall model name."""
        from .modelrelay import resolve_subcall_model
        return resolve_subcall_model(None)

    def _check_limits(self, num_calls: int = 1) -> None:
        """Check depth and call limits, raising SandboxError if exceeded."""
        if self._context.stats.depth >= self._context.config.max_depth:
            raise SandboxError(
                f"Max recursion depth ({self._context.config.max_depth}) exceeded. "
                "Consider breaking your query into smaller parts."
            )
        if self._context.stats.calls_made + num_calls > self._context.config.max_calls:
            raise SandboxError(
                f"Max sub-calls ({self._context.config.max_calls}) exceeded. "
                "Consider reducing the number of chunks or increasing --max-calls."
            )

    def llm_query(self, prompt: str, context_data: str = "") -> str:
        """Execute a recursive LLM sub-query.

        Raises:
            SandboxError: If max depth or max calls exceeded
        """
        self._check_limits(1)
        self._context.stats.calls_made += 1
        self._context.stats.depth += 1
        try:
            truncated_context = context_data[:self.MAX_CONTEXT_CHARS] if context_data else ""

            messages = [
                {"role": "system", "content": SUB_QUERY_PROMPT},
                {"role": "user", "content": f"{prompt}\n\nContext:\n{truncated_context}" if truncated_context else prompt},
            ]

            # Emit progress for sub-call via context (injectable)
            prompt_preview = prompt[:60].replace("\n", " ")
            self._context.emit_progress(f"Sub-query {self._context.stats.calls_made}/{self._context.config.max_calls}: {prompt_preview}...")

            if self._context.config.verbose:
                print(f"\n  [Sub-call {self._context.stats.calls_made}] {prompt[:80]}...", file=sys.stderr)

            response, usage = self._client.responses_create(
                messages, model=self._model, max_tokens=self.SUB_CALL_MAX_TOKENS, return_usage=True
            )

            # Track token usage
            self._context.stats.total_tokens += usage.total_tokens

            if self._context.config.verbose:
                print(f"  [Sub-call {self._context.stats.calls_made}] Response: {response[:100]}... ({usage.input_tokens}+{usage.output_tokens} tokens)", file=sys.stderr)

            return response
        finally:
            self._context.stats.depth -= 1

    def batch_llm_query(
        self,
        prompts: list[str],
        contexts: list[str] | None = None,
        *,
        max_concurrent: int | None = None,
        fail_fast: bool | None = None,
        timeout_ms: int | None = None,
    ) -> list[str]:
        """Execute multiple LLM sub-queries via ModelRelay batch API."""
        results, errors = self._run_batch(
            prompts,
            contexts,
            max_concurrent=max_concurrent,
            fail_fast=fail_fast,
            timeout_ms=timeout_ms,
        )
        if errors:
            raise BatchLLMError(
                "batch_llm_query failed for "
                f"{len(errors)}/{len(prompts)} items: "
                + json.dumps(errors, ensure_ascii=True),
                errors,
            )
        return [result if result is not None else "" for result in results]

    def llm_batch_with_errors(
        self,
        prompts: list[str],
        contexts: list[str] | None = None,
        *,
        max_concurrent: int | None = None,
        fail_fast: bool | None = None,
        timeout_ms: int | None = None,
    ) -> tuple[list[str], list[dict[str, Any]]]:
        """Execute multiple LLM sub-queries and return (results, errors)."""
        results, errors = self._run_batch(
            prompts,
            contexts,
            max_concurrent=max_concurrent,
            fail_fast=fail_fast,
            timeout_ms=timeout_ms,
        )
        normalized = [result if result is not None else "" for result in results]
        return normalized, errors

    def _run_batch(
        self,
        prompts: list[str],
        contexts: list[str] | None,
        *,
        max_concurrent: int | None,
        fail_fast: bool | None,
        timeout_ms: int | None,
    ) -> tuple[list[str | None], list[dict[str, Any]]]:
        if not prompts:
            return [], []

        contexts = contexts or [""] * len(prompts)
        if len(contexts) != len(prompts):
            raise SandboxError("prompts and contexts must have the same length")
        if len(prompts) > self.BATCH_MAX_PROMPTS:
            raise SandboxError(
                f"batch_llm_query prompt count exceeds max {self.BATCH_MAX_PROMPTS}"
            )

        self._check_limits(len(prompts))

        options: dict[str, Any] = {}
        if max_concurrent is not None:
            if not isinstance(max_concurrent, int):
                raise SandboxError("max_concurrent must be an integer")
            if max_concurrent < 1 or max_concurrent > self.BATCH_MAX_CONCURRENT:
                raise SandboxError(
                    f"max_concurrent must be between 1 and {self.BATCH_MAX_CONCURRENT}"
                )
            options["max_concurrent"] = max_concurrent
        elif len(prompts) > 1:
            options["max_concurrent"] = min(self.BATCH_MAX_CONCURRENT, len(prompts))

        if fail_fast is not None:
            if not isinstance(fail_fast, bool):
                raise SandboxError("fail_fast must be a boolean")
            options["fail_fast"] = fail_fast

        if timeout_ms is not None:
            if not isinstance(timeout_ms, int):
                raise SandboxError("timeout_ms must be an integer")
            if timeout_ms <= 0:
                raise SandboxError("timeout_ms must be greater than 0")
            options["timeout_ms"] = timeout_ms

        options_payload = options or None

        self._context.stats.calls_made += len(prompts)
        self._context.stats.depth += 1
        try:
            from .modelrelay import _build_input, _extract_text_from_response

            self._context.emit_progress(f"Processing batch of {len(prompts)} queries...")
            request_ids: list[str] = []
            request_items: list[dict[str, Any]] = []
            for idx, (prompt, context_data) in enumerate(zip(prompts, contexts)):
                truncated_context = context_data[:self.MAX_CONTEXT_CHARS] if context_data else ""
                user_prompt = (
                    f"{prompt}\n\nContext:\n{truncated_context}"
                    if truncated_context
                    else prompt
                )
                messages = [
                    {"role": "system", "content": SUB_QUERY_PROMPT},
                    {"role": "user", "content": user_prompt},
                ]
                request_id = f"subcall_{idx}"
                request_ids.append(request_id)
                request_items.append({
                    "id": request_id,
                    "model": self._model,
                    "input": _build_input(messages),
                    "max_output_tokens": self.SUB_CALL_MAX_TOKENS,
                    "temperature": 0.0,
                })

            payload = self._client.batch_responses_typed(request_items, options=options_payload)
            total_input = payload.usage.total_input_tokens
            total_output = payload.usage.total_output_tokens
            self._context.stats.total_tokens += total_input + total_output

            result_map: dict[str, str] = {}
            errors: list[dict[str, Any]] = []
            error_ids: set[str] = set()
            for item in payload.results:
                item_id = item.id
                if item.status == "success":
                    response_payload = item.response or {}
                    result_map[item_id] = _extract_text_from_response(response_payload)
                else:
                    error_payload = item.error
                    error_info = {
                        "id": item_id,
                        "status": error_payload.status if error_payload else None,
                        "code": error_payload.code if error_payload else None,
                        "message": error_payload.message if error_payload else None,
                    }
                    errors.append(error_info)
                    error_ids.add(item_id)

            for item_id in request_ids:
                if item_id not in result_map and item_id not in error_ids:
                    errors.append({
                        "id": item_id,
                        "status": None,
                        "code": "MISSING_RESULT",
                        "message": "missing batch result",
                    })

            ordered_results = [result_map.get(item_id) for item_id in request_ids]
            return ordered_results, errors
        finally:
            self._context.stats.depth -= 1

    @staticmethod
    def chunk_data(data: list[Any], chunk_size: int) -> list[list[Any]]:
        """Split data into chunks for parallel processing."""
        return [data[i:i + chunk_size] for i in range(0, len(data), chunk_size)]
