"""Multi-endpoint load balancing: one logical model, N local ports.

    from evalharness.model.pool import PooledAdapter
    from evalharness.model.adapter import resolve_adapter

    base = resolve_adapter('openai/http://127.0.0.1:8123/v1?Qwen3-8B')
    pool = PooledAdapter([resolve_adapter(f'openai/http://127.0.0.1:{p}/v1?Qwen3-8B')
                          for p in range(8123, 8131)])
    out = await pool.generate(...)   # round-robin over instances
"""

import itertools
from typing import Dict, List, Optional

from ..data.sample import ChatMessage
from .adapter import ModelAdapter
from .output import ModelOutput, Usage


class PooledAdapter(ModelAdapter):
    """Round-robin over N equivalent backend instances."""

    name = 'pool'

    def __init__(self, adapters: List[ModelAdapter]):
        if not adapters:
            raise ValueError('PooledAdapter needs at least one backend')
        super().__init__(model=adapters[0].model, api_base=adapters[0].api_base)
        self.adapters = adapters
        self._cycle = itertools.cycle(range(len(adapters)))
        self.usage = Usage()
        # request outcome counters (success rate accounting)
        self.stats = {'requests': 0, 'ok': 0, 'failed': 0, 'retried': 0}

    def _next(self) -> ModelAdapter:
        return self.adapters[next(self._cycle)]

    def request_stats(self) -> Dict[str, float]:
        """Success-rate view of all requests that passed through this pool."""
        n = self.stats['requests']
        return {
            'requests': n,
            'success_rate': self.stats['ok'] / n if n else 0.0,
            'retry_rate': self.stats['retried'] / n if n else 0.0,
            'failure_rate': self.stats['failed'] / n if n else 0.0,
        }

    async def generate(self, messages: List[ChatMessage],
                       tools: Optional[list] = None, **kw) -> ModelOutput:
        last_exc = None
        self.stats['requests'] += 1
        for _ in range(len(self.adapters)):  # try each instance once
            adapter = self._next()
            try:
                out = await adapter.generate(messages, tools=tools, **kw)
                self.usage = self.usage + out.usage
                self.stats['ok'] += 1
                if out.usage.retries:
                    self.stats['retried'] += 1
                return out
            except Exception as e:  # dead/overloaded instance -> next
                last_exc = e
                # 4xx (e.g. 400 overloaded) still worth trying ANOTHER instance:
                # one backend's state can differ from the rest
                continue
        self.stats['failed'] += 1
        raise last_exc

    async def close(self) -> None:
        for a in self.adapters:
            await a.close()


def pooled(specs: List[str]) -> PooledAdapter:
    """['openai/http://127.0.0.1:8123/v1?M', ...] -> PooledAdapter."""
    from .adapter import resolve_adapter

    return PooledAdapter([resolve_adapter(s) for s in specs])