"""Record your LLM calls and make your notebooks fast again."""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/00_core.ipynb.

# %% auto 0
__all__ = ['root_dir', 'CacheLock', 'enable_reclm']

# %% ../nbs/00_core.ipynb 3
import hashlib, httpx, json, random, time
from importlib import util as ilib_util, import_module
from unittest.mock import patch as mpatch
from fastcore.utils import *

# %% ../nbs/00_core.ipynb 31
class CacheLock:
    "Lock for our cache file."
    def __init__(self, fp):
        self.lock_fp=str(fp)+'.lock'
        self.max_retries=50
    
    def __enter__(self):
        retries=0
        while retries<self.max_retries:
            try:
                with open(self.lock_fp,'x'): return self
            except FileExistsError:
                time.sleep(0.05+random.random()*0.05)
                retries += 1
        raise TimeoutError(f"Could not acquire lock after {retries} attempts")
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        os.unlink(self.lock_fp)

# %% ../nbs/00_core.ipynb 37
def _root_dir():
    path = Path.cwd().resolve()
    while path != path.parent:
        if (path/'setup.py').exists(): return path
        path = path.parent
    return None

root_dir = _root_dir()

# %% ../nbs/00_core.ipynb 40
class _RecLMClient(httpx.Client):
    "Record and reuse your LLM calls."
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.cpth = root_dir / 'reclm.json'
        with CacheLock(self.cpth): self.cpth.touch(exist_ok=True)
        
    def _req_data(self, req): return {'method':req.method, 'url':str(req.url), 'content':req.content.decode()}
    def _hash(self, req): return hashlib.sha256(f"{json.dumps(self._req_data(req), sort_keys=True)}".encode()).hexdigest()

    def send(self, req, **kwargs):
        "Fetch `req` from the cache. If it doesn't exist, call the LLM and cache the response."
        hash = self._hash(req)
        with CacheLock(self.cpth): cache = json.loads(self.cpth.read_text() or '{}')
        if resp:= cache.get(hash): return httpx.Response(status_code=resp['status_code'], json=resp['response'], request=req)
        resp = super().send(req, **kwargs)
        cache[hash] = {'request':self._req_data(req), 'response':resp.json(), 'status_code':resp.status_code}
        with CacheLock(self.cpth): self.cpth.write_text(json.dumps(cache))
        return resp

# %% ../nbs/00_core.ipynb 44
class _RecLMCleanCacheClient(_RecLMClient):
    "Custom client used to clean the reclm cache"
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.ucpth = root_dir / 'updated_reclm.json'
        with CacheLock(self.ucpth): self.ucpth.touch(exist_ok=True)

    def send(self, req, **kwargs):
        "Save all responses to the updated cache."
        resp = super().send(req, **kwargs)
        with CacheLock(self.ucpth): ucache = json.loads(self.ucpth.read_text() or '{}')
        ucache[self._hash(req)] = {'request':self._req_data(req), 'response':resp.json(), 'status_code':resp.status_code}
        with CacheLock(self.ucpth): self.ucpth.write_text(json.dumps(ucache))
        return resp

# %% ../nbs/00_core.ipynb 46
def enable_reclm():
    "Set the OpenAI and Anthropic `http_client` to the `RecLMClient`."
    http_client = _RecLMCleanCacheClient if os.getenv('CLEAN_RECLM_CACHE') else _RecLMClient
    def _init(pkg, cls): return getattr(import_module(pkg), cls).__init__
    def _inject_http_client(oinit): return lambda *args, **kws: oinit(*args, **kws, http_client=http_client())
    sdks = {'openai': 'OpenAI', 'anthropic': 'Anthropic'}
    patches = [mpatch(f'{pkg}.{cls}.__init__', _inject_http_client(_init(pkg, cls))) for pkg, cls in sdks.items() if ilib_util.find_spec(pkg)]
    for p in patches: p.start()
    return
