from __future__ import annotations

import hashlib
import json
from dataclasses import dataclass
from typing import Mapping, Optional

from ember.types import JSONValue


@dataclass
class SUR:
    request_id: str
    org_id: str
    rbc_id: Optional[str]
    price_book_version: str
    provider_reported_usage: dict[str, int]
    canonical_usage: dict[str, int]
    correction_delta: dict[str, int]
    sur_hash: str

    def as_dict(self) -> dict[str, JSONValue]:
        return {
            "request_id": self.request_id,
            "org_id": self.org_id,
            "rbc_id": self.rbc_id,
            "price_book_version": self.price_book_version,
            "provider_reported_usage": self.provider_reported_usage,
            "canonical_usage": self.canonical_usage,
            "correction_delta": self.correction_delta,
            "sur_hash": self.sur_hash,
        }


def _hash(obj: Mapping[str, JSONValue]) -> str:
    blob = json.dumps(obj, sort_keys=True, separators=(",", ":")).encode("utf-8")
    return "sha256:" + hashlib.sha256(blob).hexdigest()


def generate_sur(
    *,
    request_id: str,
    org_id: str,
    price_book_version: str,
    provider_prompt_tokens: int,
    provider_completion_tokens: int,
    canonical_prompt_tokens: Optional[int] = None,
    canonical_completion_tokens: Optional[int] = None,
    rbc_id: Optional[str] = None,
) -> SUR:
    p_usage = {"prompt": int(provider_prompt_tokens), "completion": int(provider_completion_tokens)}
    c_prompt = int(
        canonical_prompt_tokens if canonical_prompt_tokens is not None else provider_prompt_tokens
    )
    c_completion = int(
        canonical_completion_tokens
        if canonical_completion_tokens is not None
        else provider_completion_tokens
    )
    c_usage = {"prompt": c_prompt, "completion": c_completion}
    corr = {
        "prompt": c_prompt - p_usage["prompt"],
        "completion": c_completion - p_usage["completion"],
    }
    draft = {
        "request_id": request_id,
        "org_id": org_id,
        "rbc_id": rbc_id,
        "price_book_version": price_book_version,
        "provider_reported_usage": p_usage,
        "canonical_usage": c_usage,
        "correction_delta": corr,
    }
    return SUR(
        request_id=request_id,
        org_id=org_id,
        rbc_id=rbc_id,
        price_book_version=price_book_version,
        provider_reported_usage=p_usage,
        canonical_usage=c_usage,
        correction_delta=corr,
        sur_hash=_hash(draft),
    )
