"""Batch file loader and runner for feeding multiple prompts into chains."""

from __future__ import annotations

import asyncio
import csv
import json
import logging
import time
from pathlib import Path
from typing import Any

logger = logging.getLogger(__name__)


def load_batch_file(filepath: str) -> list[str]:
    """Read a batch file and return a list of prompts.

    Supported formats:
        - .txt: each non-empty line is a prompt.
        - .csv: first column of each row (header row skipped).
        - .jsonl: the ``prompt`` field of each JSON line.

    Args:
        filepath: Path to the batch file.

    Returns:
        List of prompt strings.

    Raises:
        FileNotFoundError: If the file does not exist.
        ValueError: If the file extension is unsupported.
    """
    path = Path(filepath)
    if not path.exists():
        raise FileNotFoundError(f"Batch file not found: {filepath}")

    suffix = path.suffix.lower()

    if suffix == ".txt":
        return _load_txt(path)
    elif suffix == ".csv":
        return _load_csv(path)
    elif suffix == ".jsonl":
        return _load_jsonl(path)
    else:
        raise ValueError(f"Unsupported batch file format: {suffix}")


def _load_txt(path: Path) -> list[str]:
    """Return non-empty stripped lines from a text file."""
    text = path.read_text(encoding="utf-8")
    return [line.strip() for line in text.splitlines() if line.strip()]


def _load_csv(path: Path) -> list[str]:
    """Return first column values from a CSV, skipping the header row."""
    prompts: list[str] = []
    with path.open(newline="", encoding="utf-8") as f:
        reader = csv.reader(f)
        next(reader, None)  # skip header
        for row in reader:
            if row and row[0].strip():
                prompts.append(row[0].strip())
    return prompts


def _load_jsonl(path: Path) -> list[str]:
    """Return the 'prompt' field from each JSON line."""
    prompts: list[str] = []
    with path.open(encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            obj = json.loads(line)
            if "prompt" in obj and isinstance(obj["prompt"], str) and obj["prompt"].strip():
                prompts.append(obj["prompt"].strip())
    return prompts


# ── Batch runner ─────────────────────────────────────────────────────


async def _post_one(
    index: int,
    prompt: str,
    endpoint_url: str,
    model: str | None,
    strategy: str | None,
    semaphore: asyncio.Semaphore,
) -> dict[str, Any]:
    """Post a single prompt to *endpoint_url* and return a result dict."""
    body: dict[str, Any] = {"input": prompt}
    if model is not None:
        body["model"] = model
    if strategy is not None:
        body["strategy"] = strategy

    payload = json.dumps(body).encode("utf-8")
    headers = {"Content-Type": "application/json"}

    async with semaphore:
        start = time.monotonic()
        try:
            try:
                import aiohttp

                async with aiohttp.ClientSession() as session:
                    async with session.post(
                        endpoint_url,
                        data=payload,
                        headers=headers,
                        timeout=aiohttp.ClientTimeout(total=120),
                    ) as resp:
                        resp_body = await resp.text()
                        status = resp.status
            except ImportError:
                # Fallback: urllib in a thread pool
                status, resp_body = await _post_urllib(endpoint_url, payload, headers)

            elapsed = (time.monotonic() - start) * 1000

            try:
                output = json.loads(resp_body)
            except (json.JSONDecodeError, TypeError):
                output = resp_body

            return {
                "index": index,
                "input": prompt,
                "output": output,
                "status": status,
                "latency_ms": round(elapsed, 1),
            }
        except Exception as exc:
            elapsed = (time.monotonic() - start) * 1000
            logger.warning("Batch item %d failed: %s", index, exc)
            return {
                "index": index,
                "input": prompt,
                "output": str(exc),
                "status": 0,
                "latency_ms": round(elapsed, 1),
            }


async def _post_urllib(
    url: str, payload: bytes, headers: dict[str, str]
) -> tuple[int, str]:
    """Blocking urllib POST wrapped in a thread-pool executor."""
    import urllib.error
    import urllib.request

    def _do() -> tuple[int, str]:
        req = urllib.request.Request(url, data=payload, headers=headers, method="POST")
        try:
            with urllib.request.urlopen(req, timeout=120) as resp:
                return resp.status, resp.read().decode("utf-8", errors="replace")
        except urllib.error.HTTPError as exc:
            return exc.code, exc.read().decode("utf-8", errors="replace")

    loop = asyncio.get_running_loop()
    return await loop.run_in_executor(None, _do)


async def run_batch(
    prompts: list[str],
    endpoint_url: str,
    model: str | None = None,
    strategy: str | None = None,
    max_parallel: int = 4,
) -> list[dict]:
    """Post each prompt to *endpoint_url* concurrently and collect results.

    Args:
        prompts: List of prompt strings to send.
        endpoint_url: URL to POST each prompt to (e.g. runtime /chat or /chain).
        model: Optional model name to include in the request body.
        strategy: Optional chain strategy to include in the request body.
        max_parallel: Maximum number of concurrent requests (default 4).

    Returns:
        List of dicts, one per prompt, sorted by index::

            {index, input, output, status, latency_ms}
    """
    if max_parallel < 1:
        max_parallel = 1

    semaphore = asyncio.Semaphore(max_parallel)
    tasks = [
        _post_one(i, prompt, endpoint_url, model, strategy, semaphore)
        for i, prompt in enumerate(prompts)
    ]
    results = await asyncio.gather(*tasks)
    return sorted(results, key=lambda r: r["index"])


def write_batch_results(results: list[dict], filepath: str, fmt: str) -> None:
    """Write batch results to a file in the specified format.

    Args:
        results: List of result dicts to write.
        filepath: Destination file path.
        fmt: Output format — ``'json'``, ``'csv'``, or ``'jsonl'``.

    Raises:
        ValueError: If *fmt* is not one of the supported formats.
    """
    if fmt not in ("json", "csv", "jsonl"):
        raise ValueError(f"Unsupported output format: {fmt}")

    path = Path(filepath)

    if fmt == "json":
        path.write_text(json.dumps(results, indent=2, default=str), encoding="utf-8")
    elif fmt == "csv":
        if not results:
            path.write_text("", encoding="utf-8")
            return
        headers = list(results[0].keys())
        with path.open("w", newline="", encoding="utf-8") as f:
            writer = csv.DictWriter(f, fieldnames=headers)
            writer.writeheader()
            writer.writerows(results)
    elif fmt == "jsonl":
        with path.open("w", encoding="utf-8") as f:
            for item in results:
                f.write(json.dumps(item, default=str) + "\n")


def batch_error_report(results: list[dict]) -> dict:
    """Analyze failed results from a batch run and return an error summary.

    A result is considered failed if its ``status`` is not in the 200-299 range.

    Args:
        results: List of result dicts as returned by :func:`run_batch`.

    Returns:
        Dict with ``total_errors``, ``errors_by_type``, ``failed_indices``,
        and a suggested ``retry_command``.
    """
    errors_by_type: dict[str, int] = {}
    failed_indices: list[int] = []

    for r in results:
        status = r.get("status", 0)
        if isinstance(status, int) and 200 <= status <= 299:
            continue
        idx = r.get("index", 0)
        failed_indices.append(idx)
        error_type = f"http_{status}" if isinstance(status, int) and status > 0 else "connection_error"
        errors_by_type[error_type] = errors_by_type.get(error_type, 0) + 1

    if failed_indices:
        indices_str = ",".join(str(i) for i in failed_indices)
        retry_command = f"brynq-runtime batch --retry-indices {indices_str}"
    else:
        retry_command = ""

    return {
        "total_errors": len(failed_indices),
        "errors_by_type": errors_by_type,
        "failed_indices": failed_indices,
        "retry_command": retry_command,
    }


# ── Batch progress tracker ───────────────────────────────────────────


class BatchProgress:
    """Track progress of a batch run."""

    def __init__(self, total: int) -> None:
        self.total = total
        self.completed = 0
        self.failed = 0
        self._start = time.monotonic()

    def update(self, completed: int, failed: int) -> None:
        """Update progress counters."""
        self.completed = completed
        self.failed = failed

    def summary(self) -> dict:
        """Return a snapshot of current progress."""
        elapsed = time.monotonic() - self._start
        done = self.completed + self.failed
        if done > 0 and done < self.total:
            eta = elapsed * (self.total - done) / done
        else:
            eta = 0.0
        percent = (done / self.total * 100) if self.total > 0 else 0.0
        return {
            "total": self.total,
            "completed": self.completed,
            "failed": self.failed,
            "percent": round(percent, 1),
            "elapsed_seconds": round(elapsed, 2),
            "eta_seconds": round(eta, 2),
        }
