import time
from typing import Dict, Tuple

class TokenCache:
    def __init__(self):
        self.cache: Dict[str, Tuple[str, int]] = {}

    def get(self, plan_hash: str) -> str | None:
        if plan_hash in self.cache:
            token, expires_at = self.cache[plan_hash]
            # Return token if valid for at least 60 more seconds
            if time.time() < expires_at - 60:
                return token
        return None

    def set(self, plan_hash: str, token: str, expires_at: int):
        self.cache[plan_hash] = (token, expires_at)

    def clear_expired(self):
        now = time.time()
        self.cache = {
            k: v for k, v in self.cache.items()
            if v[1] > now
        }

# Usage
token_cache = TokenCache()

def get_token_cached(client, llm, prompt):
    captured = client.capture_plan(llm=llm, prompt=prompt)
    plan_hash = captured.plan_hash

    # Try cache first
    token = token_cache.get(plan_hash)
    if token:
        return token

    # Get new token
    response = client.get_intent_token(captured)
    token_cache.set(plan_hash, response["token"], response["expires_at"])

    return response["token"]

import { IntentToken, ArmorIQClient, PlanCapture } from '@armoriq/sdk';

class TokenCache {
  private cache: Map<string, IntentToken> = new Map();

  get(planHash: string): IntentToken | undefined {
    const token = this.cache.get(planHash);
    if (token) {
      // Return token if valid for at least 60 more seconds
      if (IntentToken.timeUntilExpiry(token) > 60) {
        return token;
      }
      // Token expired or expiring soon, remove from cache
      this.cache.delete(planHash);
    }
    return undefined;
  }

  set(planHash: string, token: IntentToken): void {
    this.cache.set(planHash, token);
  }

  clearExpired(): void {
    for (const [hash, token] of this.cache.entries()) {
      if (IntentToken.isExpired(token)) {
        this.cache.delete(hash);
      }
    }
  }
}

// Usage
const tokenCache = new TokenCache();

async function getTokenCached(
  client: ArmorIQClient,
  llm: string,
  prompt: string,
  plan: Record<string, any>
): Promise<IntentToken> {
  const captured = client.capturePlan(llm, prompt, plan);
  const planHash = captured.plan?.hash || JSON.stringify(captured.plan);

  // Try cache first
  const cachedToken = tokenCache.get(planHash);
  if (cachedToken) {
    console.log('Using cached token');
    return cachedToken;
  }

  // Get new token
  const token = await client.getIntentToken(captured);
  tokenCache.set(planHash, token);

  return token;
}