"""Core: ApiConfig, Task, Sample, complete(), run_task()."""

from __future__ import annotations

import asyncio
import base64
import hashlib
import logging
import math
import time
from collections import Counter
from collections.abc import Callable
from dataclasses import dataclass, field
from io import BytesIO
from pathlib import Path
from typing import Any

import datasets
import datasets.config as ds_config
import httpx
from datasets import Dataset, DownloadMode
from huggingface_hub.constants import HF_HOME, HF_HUB_CACHE
from PIL import Image
from tqdm.asyncio import tqdm_asyncio

from nano_eval._types import (
    DEFAULT_MAX_CONCURRENT,
    DEFAULT_REQUEST_TIMEOUT,
    Metrics,
    TaskResult,
)

logger = logging.getLogger("nano_eval.core")


@dataclass(frozen=True)
class Prompt:
    """Evaluation prompt with optional images for multimodal tasks."""

    text: str | list[dict[str, str]]
    images: list[Any] = field(default_factory=list)

    def to_messages(self) -> list[dict[str, Any]]:
        if self.images:
            assert isinstance(self.text, str)
            content: list[dict[str, Any]] = [
                {
                    "type": "image_url",
                    "image_url": {"url": f"data:image/png;base64,{b64}"},
                }
                for img in self.images
                if (b64 := _encode_image(img))
            ]
            content.append({"type": "text", "text": self.text})
            return [{"role": "user", "content": content}]
        if isinstance(self.text, list):
            return self.text
        return [{"role": "user", "content": self.text}]

    def text_repr(self) -> str:
        if isinstance(self.text, list):
            return "\n".join(f"{m['role']}: {m['content']}" for m in self.text)
        return self.text


@dataclass
class Sample:
    """Prompt + expected target."""

    prompt: Prompt
    target: str


@dataclass(frozen=True)
class Task:
    """Dataset config + scoring function."""

    name: str
    dataset: str
    revision: str
    splits: list[str]
    extract: Callable[[dict[str, Any]], Sample]
    score: Callable[[str, str], float]
    config_name: str | None = None

    def load_samples(
        self, max_samples: int | None = None, seed: int | None = None
    ) -> list[Sample]:
        """Load samples from HuggingFace dataset across splits."""
        # TODO Upstream fix. HF datasets logging is too noisy
        datasets.utils.logging.set_verbosity_error()

        cache_path = (
            Path(HF_HUB_CACHE)
            / f"datasets--{self.dataset.replace('/', '--')}"
            / "snapshots"
            / self.revision
        )
        cached = cache_path.is_dir()
        logger.info(
            f"Cache {'hit' if cached else 'miss'} for {self.dataset}, HF_HOME={HF_HOME}"
        )
        old_offline = ds_config.HF_HUB_OFFLINE
        if cached:
            ds_config.HF_HUB_OFFLINE = True

        try:
            result: list[Sample] = []
            for split in self.splits:
                if max_samples is not None and len(result) >= max_samples:
                    break
                ds = datasets.load_dataset(
                    self.dataset,
                    name=self.config_name,
                    split=split,
                    revision=self.revision,
                    download_mode=DownloadMode.REUSE_DATASET_IF_EXISTS,
                )
                assert isinstance(ds, Dataset)
                if seed is not None:
                    ds = ds.shuffle(seed=seed)
                if max_samples is not None:
                    ds = ds.select(range(min(max_samples - len(result), len(ds))))
                result.extend(self.extract(doc) for doc in ds)
            return result
        finally:
            ds_config.HF_HUB_OFFLINE = old_offline


@dataclass
class ApiConfig:
    """API endpoint configuration."""

    url: str
    model: str
    api_key: str = ""
    max_concurrent: int = DEFAULT_MAX_CONCURRENT
    timeout: int = DEFAULT_REQUEST_TIMEOUT
    gen_kwargs: dict[str, Any] = field(default_factory=dict)


async def _request(
    client: httpx.AsyncClient,
    url: str,
    payload: dict[str, Any],
) -> dict[str, Any] | None:
    t0 = time.perf_counter()
    try:
        resp = await client.post(url, json=payload)
    except httpx.TimeoutException:
        logger.warning(f"Request timed out after {time.perf_counter() - t0:.1f}s")
        return None
    except httpx.HTTPError as exc:
        logger.warning(f"Request failed: {exc}")
        return None
    if resp.is_success:
        data = resp.json()
        choice = data["choices"][0]
        return {
            "answer": choice["message"]["content"],
            "stop_reason": choice["finish_reason"],
            "input_tokens": data["usage"]["prompt_tokens"],
            "output_tokens": data["usage"]["completion_tokens"],
            "duration_seconds": time.perf_counter() - t0,
        }
    logger.warning(f"Request failed ({resp.status_code}): {resp.text}")
    return None


async def complete(
    prompts: list[Prompt],
    config: ApiConfig,
    progress_desc: str = "Running evals",
) -> list[dict[str, Any] | None]:
    """Run batch of chat completions with concurrency control."""
    headers: dict[str, str] = {"Content-Type": "application/json"}
    if config.api_key:
        headers["Authorization"] = f"Bearer {config.api_key}"

    async with httpx.AsyncClient(
        limits=httpx.Limits(max_connections=config.max_concurrent),
        timeout=httpx.Timeout(config.timeout),
        headers=headers,
        trust_env=True,
    ) as client:
        tasks: list[asyncio.Task[dict[str, Any] | None]] = []
        for prompt in prompts:
            payload: dict[str, Any] = {
                "model": config.model,
                "messages": prompt.to_messages(),
                **config.gen_kwargs,
            }
            tasks.append(asyncio.create_task(_request(client, config.url, payload)))

        try:
            return list(
                await tqdm_asyncio.gather(*tasks, desc=progress_desc, leave=False)
            )
        except BaseException:
            for task in tasks:
                task.cancel()
            await asyncio.gather(*tasks, return_exceptions=True)
            raise


def _encode_image(image: Any) -> str:
    """Encode PIL image to base64, or pass through base64 string."""
    if isinstance(image, str):
        if image.startswith("http"):
            raise ValueError("Remote image URLs are not supported.")
        return image

    if isinstance(image, Image.Image):
        if image.mode not in ("RGB", "L"):
            image = image.convert("RGB")
        buf = BytesIO()
        image.save(buf, format="PNG")
        return base64.b64encode(buf.getvalue()).decode()

    raise TypeError(f"Unsupported image type: {type(image).__name__}")


def compute_samples_hash(samples: list[Sample]) -> str:
    """SHA256 hash of all samples for reproducibility tracking."""
    hasher = hashlib.sha256()
    for s in samples:
        hasher.update(s.prompt.text_repr().encode())
        for img in s.prompt.images:
            hasher.update(_encode_image(img).encode())
        hasher.update(s.target.encode())
    return hasher.hexdigest()


async def run_task(
    task: Task,
    config: ApiConfig,
    modality: str,
    max_samples: int | None = None,
    seed: int | None = None,
) -> tuple[TaskResult, list[dict[str, Any]]]:
    """Evaluate a task: load samples, run inference, compute scores."""
    samples = task.load_samples(max_samples, seed)
    samples_hash = compute_samples_hash(samples)
    prompts = [s.prompt for s in samples]

    logger.info(
        f"Starting {modality} ({task.name}) eval: "
        f"{len(samples)} samples, up to {config.max_concurrent} concurrent requests"
    )
    t0 = time.perf_counter()
    raw_responses = await complete(prompts, config, f"Running {modality} eval")
    elapsed = time.perf_counter() - t0

    timeout_count = sum(1 for r in raw_responses if r is None)
    valid = [(s, r) for s, r in zip(samples, raw_responses) if r is not None]
    valid_samples = [s for s, _ in valid]
    responses = [r for _, r in valid]

    reason_counts = Counter(r["stop_reason"] for r in responses)
    non_stop = len(responses) - reason_counts.get("stop", 0)
    if non_stop:
        logger.warning(
            f"{non_stop} response(s) did not finish with 'stop'. "
            f"Reasons: {dict(reason_counts)}"
        )

    scores = [task.score(r["answer"], s.target) for s, r in valid]
    n = len(valid)

    if logger.isEnabledFor(logging.DEBUG):
        for i, (s, r, score) in enumerate(zip(valid_samples, responses, scores)):
            status = "PASS" if score == 1.0 else "FAIL"
            logger.debug(
                f"[{i + 1}/{n} {status}] target={s.target!r} got={r['answer']!r}"
            )

    accuracy = sum(scores) / n if n else 0.0
    stderr = math.sqrt(accuracy * (1 - accuracy) / n) if n > 0 else 0.0
    logger.debug(
        f"{task.name}: accuracy={accuracy:.4f}+/-{stderr:.4f} ({elapsed:.2f}s)"
    )

    if timeout_count:
        logger.error(
            f"{timeout_count}/{len(raw_responses)} request(s) timed out "
            f"and were excluded from results"
        )

    request_logs = [
        {
            "request_id": i,
            "target": s.target,
            "prompt": s.prompt.text_repr(),
            "response": r["answer"],
            "score": score,
            "stop_reason": r["stop_reason"],
            "input_tokens": r["input_tokens"],
            "output_tokens": r["output_tokens"],
            "duration_seconds": r["duration_seconds"],
        }
        for i, (s, r, score) in enumerate(zip(valid_samples, responses, scores))
    ]

    total_input = sum(r["input_tokens"] for r in responses)
    total_output = sum(r["output_tokens"] for r in responses)
    sum_of_request_durations_seconds = sum(r["duration_seconds"] for r in responses)

    result = TaskResult(
        elapsed_seconds=elapsed,
        metrics=Metrics(accuracy=accuracy, accuracy_stderr=stderr),
        num_samples=n,
        samples_hash=samples_hash,
        task=task.name,
        modality=modality,
        total_input_tokens=total_input,
        total_output_tokens=total_output,
        per_request_tokens_per_second=(total_input + total_output)
        / sum_of_request_durations_seconds
        if sum_of_request_durations_seconds
        else 0.0,
    )
    return result, request_logs
