"""Insight ingestion client for traces and scores."""

from __future__ import annotations

import base64
import json
import os
import uuid
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import Any, Literal
from urllib import error, request


def _utc_now_iso() -> str:
    return datetime.now(timezone.utc).isoformat(timespec="milliseconds").replace(
        "+00:00", "Z"
    )


@dataclass
class InsightIngestionConfig:
    """Configuration for Insight ingestion."""

    host: str
    public_key: str
    secret_key: str
    timeout_seconds: int = 30
    # Keep default aligned with importer implementation in Insight web app.
    max_batch_size: int = 200


class InsightIngestionClient:
    """Client for sending trace and score data to Insight."""

    def __init__(self, config: InsightIngestionConfig):
        self._host = config.host.rstrip("/")
        self._timeout_seconds = config.timeout_seconds
        self._max_batch_size = config.max_batch_size

        token = f"{config.public_key}:{config.secret_key}".encode("utf-8")
        self._auth_header = f"Basic {base64.b64encode(token).decode('utf-8')}"

    @classmethod
    def from_env(cls) -> "InsightIngestionClient":
        """Build a client from INSIGHT_* environment variables."""
        host = os.getenv("INSIGHT_HOST")
        public_key = os.getenv("INSIGHT_PUBLIC_KEY")
        secret_key = os.getenv("INSIGHT_SECRET_KEY")

        missing = [
            name
            for name, value in [
                ("INSIGHT_HOST", host),
                ("INSIGHT_PUBLIC_KEY", public_key),
                ("INSIGHT_SECRET_KEY", secret_key),
            ]
            if not value
        ]
        if missing:
            raise ValueError(
                "Missing required environment variables for Insight ingestion client: "
                + ", ".join(missing)
            )

        return cls(
            InsightIngestionConfig(
                host=host,
                public_key=public_key,
                secret_key=secret_key,
            )
        )

    def _post_json(self, path: str, payload: dict[str, Any]) -> dict[str, Any]:
        body = json.dumps(payload).encode("utf-8")
        req = request.Request(
            f"{self._host}{path}",
            data=body,
            method="POST",
            headers={
                "Content-Type": "application/json",
                "Authorization": self._auth_header,
            },
        )
        try:
            with request.urlopen(req, timeout=self._timeout_seconds) as response:
                raw = response.read().decode("utf-8")
                return json.loads(raw) if raw else {}
        except error.HTTPError as exc:
            raw = exc.read().decode("utf-8", errors="ignore")
            raise RuntimeError(
                f"Insight API error {exc.code} on {path}: {raw or exc.reason}"
            ) from exc

    def ingest_trace_event(
        self,
        *,
        trace_id: str,
        name: str,
        input_data: Any,
        output_data: Any | None = None,
        session_id: str | None = None,
        user_id: str | None = None,
        tags: list[str] | None = None,
        metadata: dict[str, Any] | None = None,
        environment: str = "default",
        trace_timestamp: str | None = None,
    ) -> dict[str, Any]:
        """Send one trace-create event to Insight."""
        timestamp = trace_timestamp or _utc_now_iso()
        event = {
            "id": str(uuid.uuid4()),
            "type": "trace-create",
            "timestamp": timestamp,
            "body": {
                "id": trace_id,
                "timestamp": timestamp,
                "name": name,
                "input": input_data,
                "output": output_data,
                "sessionId": session_id,
                "userId": user_id,
                "tags": tags or [],
                "metadata": metadata or {},
                "environment": environment,
            },
        }
        return self._post_json("/api/public/ingestion", {"batch": [event]})

    def send_score_api(self, score: dict[str, Any]) -> dict[str, Any]:
        """Send one score directly via /api/public/scores."""
        return self._post_json("/api/public/scores", score)

    def send_scores(
        self,
        scores: list[dict[str, Any]],
        *,
        mode: Literal["ingestion-batch", "scores-api"] = "ingestion-batch",
        batch_size: int | None = None,
    ) -> list[dict[str, Any]]:
        """Send many scores efficiently."""
        if not scores:
            return []

        if mode == "scores-api":
            return [self.send_score_api(score) for score in scores]

        effective_batch_size = batch_size or self._max_batch_size
        if effective_batch_size <= 0:
            raise ValueError("batch_size must be > 0")

        responses: list[dict[str, Any]] = []
        for i in range(0, len(scores), effective_batch_size):
            subset = scores[i : i + effective_batch_size]
            events = []
            now_ts = _utc_now_iso()
            for score in subset:
                body = dict(score)
                body.setdefault("id", str(uuid.uuid4()))
                body.setdefault("environment", "default")
                events.append(
                    {
                        "id": str(uuid.uuid4()),
                        "type": "score-create",
                        "timestamp": now_ts,
                        "body": body,
                    }
                )
            responses.append(self._post_json("/api/public/ingestion", {"batch": events}))
        return responses


__all__ = [
    "InsightIngestionClient",
    "InsightIngestionConfig",
]
