api/utils.py

import asyncio
import json
import random
import re
import string
import time
import uuid
from datetime import datetime, timezone, timedelta
from typing import Any, Dict, List, Optional

import boto3
import httpx
import tiktoken
import platform
import hashlib
from fastapi import HTTPException

from api.config import (
    MODEL_MAPPING,
    get_headers_api_chat,
    get_headers_chat,
    BASE_URL,
    AGENT_MODE,
    TRENDING_AGENT_MODE,
    MODEL_PREFIXES
)
from api.logger import setup_logger
from api.models import ChatRequest
from api.validate import getHid  # Import the asynchronous getHid function

logger = setup_logger(__name__)

# ---------------------------------------------
#         CLOUDFLARE R2 CONFIGURATION
# ---------------------------------------------
R2_ACCESS_KEY_ID = "df9c9eb87e850a8eb27afd3968077b42"
R2_SECRET_ACCESS_KEY = "14b08b0855263bb63d2618da3a6537e1b0446d89d51da03a568620b1e5342ea8"
R2_ENDPOINT_URL = "https://f2f92ac53fae792c4155f6e93a514989.r2.cloudflarestorage.com"
R2_BUCKET_NAME = "snapzion"
R2_REPLACED_URLS_KEY = "snapzion.txt"

s3 = boto3.client(
    "s3",
    endpoint_url=R2_ENDPOINT_URL,
    aws_access_key_id=R2_ACCESS_KEY_ID,
    aws_secret_access_key=R2_SECRET_ACCESS_KEY,
)

BLOCKED_MESSAGE = (
    "Generated by BLACKBOX.AI, try unlimited chat https://www.blackbox.ai "
    "and for API requests replace https://www.blackbox.ai with https://api.blackbox.ai"
)

# ---------------------------------------------
#   RANDOM USER-DATA & SESSION GENERATION
# ---------------------------------------------
def get_random_name_email_customer():
    first_names = ["Aliace", "B21ob", "Car232ol", "Daavid", "Evewwlyn", "Fraank", "Grssace", "Hefctor", "Ivgy", "Jackdie"]
    last_names = ["Smilth", "Johnkson", "Dajvis", "Mihller", "Thomgpson", "Garwcia", "Broawn", "Wilfson", "Maartin", "Clarak"]

    random_name = f"{random.choice(first_names)} {random.choice(last_names)}"
    email_username = ''.join(random.choices(string.ascii_lowercase + string.digits, k=8))
    random_email = f"[email protected]"
    suffix_length = len("Rldf7IKdNhdhiw")
    suffix_chars = string.ascii_letters + string.digits
    random_suffix = ''.join(random.choice(suffix_chars) for _ in range(suffix_length))
    random_customer_id = f"cus_{random_suffix}"

    return random_name, random_email, random_customer_id

def generate_session(email: str, id_length: int = 21, days_ahead: int = 365) -> dict:
    numeric_id = ''.join(random.choice('0123456789') for _ in range(id_length))
    future_date = datetime.now(timezone.utc) + timedelta(days=days_ahead)
    expiry = future_date.isoformat(timespec='milliseconds').replace('+00:00', 'Z')

    chars = string.ascii_letters + string.digits + "-"
    random_img_id = ''.join(random.choice(chars) for _ in range(48))
    image_url = f"https://lh3.googleusercontent.com/a/ACg8oc{random_img_id}=s96-c"

    return {
        "user": {
            "name": "SNAPZION",
            "email": email,
            "image": image_url,
            "id": numeric_id
        },
        "expires": expiry,
        "isNewUser": False
    }

def generate_session_data() -> dict:
    _, email, _ = get_random_name_email_customer()
    session_data = generate_session(email)
    logger.info(f"Using generated session with email {email}")
    return session_data

# ---------------------------------------------
#         HELPER FUNCTIONS
# ---------------------------------------------
def generate_system_fingerprint() -> str:
    raw = f"{platform.node()}-{time.time()}-{uuid.uuid4()}"
    short_hash = hashlib.md5(raw.encode()).hexdigest()[:12]
    return f"fp_{short_hash}"

def get_last_user_prompt(messages: List[Any]) -> str:
    for msg in reversed(messages):
        if msg.role == "user":
            if isinstance(msg.content, str):
                return msg.content.strip()
            elif isinstance(msg.content, list):
                for item in msg.content:
                    if item.get("type") == "text":
                        return item.get("text", "").strip()
    return ""

def upload_replaced_urls_to_r2(urls: List[str], alt_text: str = "") -> None:
    if not urls:
        logger.info("No replaced or final URLs to store.")
        return
    existing = ""
    try:
        resp = s3.get_object(Bucket=R2_BUCKET_NAME, Key=R2_REPLACED_URLS_KEY)
        existing = resp["Body"].read().decode()
    except s3.exceptions.NoSuchKey:
        pass
    except Exception as e:
        logger.error(f"Error reading {R2_REPLACED_URLS_KEY}: {e}")

    markdown = "\n".join(f"![{alt_text}]({u})" for u in urls)
    content = f"{existing}\n{markdown}" if existing.strip() else markdown

    try:
        s3.put_object(
            Bucket=R2_BUCKET_NAME,
            Key=R2_REPLACED_URLS_KEY,
            Body=content.encode(),
            ContentType="text/plain",
        )
        logger.info(f"Appended {len(urls)} URLs to {R2_REPLACED_URLS_KEY}.")
    except Exception as e:
        logger.error(f"Error writing {R2_REPLACED_URLS_KEY}: {e}")

def calculate_tokens(text: str, model: str) -> int:
    try:
        enc = tiktoken.encoding_for_model(model)
        return len(enc.encode(text))
    except Exception:
        return len(text.split())

def create_chat_completion_data(
    content: str,
    model: str,
    timestamp: int,
    request_id: str,
    system_fingerprint: str,
    prompt_tokens: int = 0,
    completion_tokens: int = 0,
    finish_reason: Optional[str] = None,
) -> Dict[str, Any]:
    usage = None
    if finish_reason == "stop":
        usage = {
            "prompt_tokens": prompt_tokens,
            "completion_tokens": completion_tokens,
            "total_tokens": prompt_tokens + completion_tokens,
        }
    return {
        "id": request_id,
        "object": "chat.completion.chunk",
        "created": timestamp,
        "model": model,
        "system_fingerprint": system_fingerprint,
        "choices": [{"index": 0, "delta": {"content": content, "role": "assistant"}, "finish_reason": finish_reason}],
        "usage": usage,
    }

def message_to_dict(message, model_prefix: Optional[str] = None):
    """
    Convert a ChatRequest message to a dict for the request payload.
    Supports up to three images with type-based structure and sends multiple formats.
    Prepends model_prefix to text content if specified.
    """
    content = ""
    images_data = []
    image_urls = []

    if isinstance(message.content, list):
        for item in message.content:
            if item.get("type") == "text":
                content = item.get("text", "").strip()
            elif item.get("type") == "image_url" and len(images_data) < 3:
                url = item["image_url"].get("url", "")
                if url:
                    path = f"MultipleFiles/{uuid.uuid4().hex}.jpg"
                    images_data.append({"filePath": path, "contents": url})
                    image_urls.append({"image_url": {"url": url}})
    elif isinstance(message.content, str):
        content = message.content.strip()

    if model_prefix and content:
        content = f"{model_prefix} {content}"

    base = {"role": message.role, "content": content}
    if images_data:
        base["data"] = {
            "imageBase64": images_data[0]["contents"],
            "fileText": "",
            "title": "snapshot",
            "imagesData": images_data
        }
        # embed any extra URLs in content list
        for extra in image_urls[1:]:
            base.setdefault("content", []).append(extra)

    return base

def strip_model_prefix(content: str, model_prefix: Optional[str] = None) -> str:
    if model_prefix and content.startswith(model_prefix):
        return content[len(model_prefix):].strip()
    return content

# ---------------------------------------------
#    STREAMING RESPONSE HANDLER
# ---------------------------------------------
async def process_streaming_response(request: ChatRequest):
    system_fp = generate_system_fingerprint()
    request_id = f"chatcmpl-{uuid.uuid4()}"
    logger.info(f"Processing (stream) {request_id} - Model: {request.model}")

    agent_mode = AGENT_MODE.get(request.model, {})
    trending_mode = TRENDING_AGENT_MODE.get(request.model, {})
    prefix = MODEL_PREFIXES.get(request.model, "")

    headers_api = get_headers_api_chat(BASE_URL)
    if request.model == "o1-preview":
        await asyncio.sleep(random.randint(1, 60))

    h = await getHid()
    if not h:
        raise HTTPException(status_code=500, detail="Missing h-value.")

    msgs = [message_to_dict(m, prefix) for m in request.messages]

    json_data = {
        "agentMode": agent_mode,
        "clickedAnswer2": False,
        "clickedAnswer3": False,
        "reasoningMode": False,
        "clickedForceWebSearch": False,
        "codeInterpreterMode": False,
        "codeModelMode": True,
        "githubToken": "",
        "deepSearchMode": False,
        "domains": None,
        "id": request_id,
        "imageGenerationMode": False,
        "isChromeExt": False,
        "isMicMode": False,
        "isPremium": True,
        "isMemoryEnabled": False,
        "maxTokens": request.max_tokens,
        "messages": msgs,
        "mobileClient": False,
        "playgroundTemperature": request.temperature,
        "playgroundTopP": request.top_p,
        "previewToken": None,
        "trendingAgentMode": trending_mode,
        "userId": None,
        "userSelectedModel": MODEL_MAPPING.get(request.model, request.model),
        "userSystemPrompt": None,
        "validated": h,
        "visitFromDelta": False,
        "webSearchModePrompt": False,
        "vscodeClient": False,
        "designerMode": False,
        "workspaceId": "",
        "beastMode": False,
        "customProfile": {"name": "", "occupation": "", "traits": [], "additionalInfo": "", "enableNewChats": False},
        "webSearchModeOption": {"autoMode": False, "webMode": False, "offlineMode": True},
        "session": generate_session_data(),
    }

    prompt_tokens = sum(
        calculate_tokens(m.get("content", ""), request.model) +
        sum(calculate_tokens(img["contents"], request.model) for img in m.get("data", {}).get("imagesData", []))
        for m in msgs
    )

    completion_tokens = 0
    final_links: List[str] = []

    async with httpx.AsyncClient() as client:
        try:
            async with client.stream("POST", f"{BASE_URL}/api/chat", headers=headers_api, json=json_data, timeout=100) as resp:
                resp.raise_for_status()
                async for chunk in resp.aiter_text():
                    if not chunk:
                        continue
                    if chunk.startswith("$@$v=undefined-rv1$@$"):
                        chunk = chunk[21:]
                    if BLOCKED_MESSAGE in chunk:
                        chunk = chunk.replace(BLOCKED_MESSAGE, "").strip()
                        if not chunk:
                            continue
                    chunk = chunk.replace("https://storage.googleapis.com", "https://cdn.snapzion.com")
                    links = re.findall(r"(https://cdn\.snapzion\.com[^\s\)]+)", chunk)
                    final_links.extend(links)
                    clean = strip_model_prefix(chunk, prefix)
                    completion_tokens += calculate_tokens(clean, request.model)
                    ts = int(datetime.now().timestamp())
                    yield "data: " + json.dumps(
                        create_chat_completion_data(clean, request.model, ts, request_id, system_fp, prompt_tokens, completion_tokens)
                    ) + "\n\n"

                ts = int(datetime.now().timestamp())
                yield "data: " + json.dumps(
                    create_chat_completion_data("", request.model, ts, request_id, system_fp, prompt_tokens, completion_tokens, "stop")
                ) + "\n\n"
                yield "data: [DONE]\n\n"

        except Exception as e:
            logger.error(f"Stream error {request_id}: {e}")

    upload_replaced_urls_to_r2(final_links, alt_text=get_last_user_prompt(request.messages))

# ---------------------------------------------
#     NON-STREAMING RESPONSE HANDLER
# ---------------------------------------------
async def process_non_streaming_response(request: ChatRequest):
    system_fp = generate_system_fingerprint()
    request_id = f"chatcmpl-{uuid.uuid4()}"
    logger.info(f"Processing (non-stream) {request_id} - Model: {request.model}")

    agent_mode = AGENT_MODE.get(request.model, {})
    trending_mode = TRENDING_AGENT_MODE.get(request.model, {})
    prefix = MODEL_PREFIXES.get(request.model, "")

    headers_api = get_headers_api_chat(BASE_URL)
    if request.model == "o1-preview":
        await asyncio.sleep(random.randint(20, 60))

    h = await getHid()
    if not h:
        raise HTTPException(status_code=500, detail="Missing h-value.")

    msgs = [message_to_dict(m, prefix) for m in request.messages]

    json_data = {
        "agentMode": agent_mode,
        "clickedAnswer2": False,
        "clickedAnswer3": False,
        "reasoningMode": False,
        "clickedForceWebSearch": False,
        "codeInterpreterMode": False,
        "codeModelMode": True,
        "githubToken": "",
        "deepSearchMode": False,
        "domains": None,
        "id": request_id,
        "imageGenerationMode": False,
        "isChromeExt": False,
        "isMicMode": False,
        "isPremium": True,
        "isMemoryEnabled": False,
        "maxTokens": request.max_tokens,
        "messages": msgs,
        "mobileClient": False,
        "playgroundTemperature": request.temperature,
        "playgroundTopP": request.top_p,
        "previewToken": None,
        "trendingAgentMode": trending_mode,
        "userId": None,
        "userSelectedModel": MODEL_MAPPING.get(request.model, request.model),
        "userSystemPrompt": None,
        "validated": h,
        "visitFromDelta": False,
        "webSearchModePrompt": False,
        "vscodeClient": False,
        "designerMode": False,
        "workspaceId": "",
        "beastMode": False,
        "customProfile": {"name": "", "occupation": "", "traits": [], "additionalInfo": "", "enableNewChats": False},
        "webSearchModeOption": {"autoMode": False, "webMode": False, "offlineMode": True},
        "session": generate_session_data(),
    }

    prompt_tokens = sum(
        calculate_tokens(m.get("content", ""), request.model) +
        sum(calculate_tokens(img["contents"], request.model) for img in m.get("data", {}).get("imagesData", []))
        for m in msgs
    )

    full_resp = ""
    final_links: List[str] = []

    async with httpx.AsyncClient() as client:
        try:
            resp = await client.post(f"{BASE_URL}/api/chat", headers=headers_api, json=json_data)
            resp.raise_for_status()
            full_resp = resp.text
        except Exception as e:
            return {
                "id": request_id,
                "object": "chat.completion",
                "created": int(datetime.now().timestamp()),
                "model": request.model,
                "system_fingerprint": system_fp,
                "choices": [{"index": 0, "message": {"role": "assistant", "content": str(e)}, "finish_reason": "error"}],
                "usage": {"prompt_tokens": prompt_tokens, "completion_tokens": 0, "total_tokens": prompt_tokens},
            }

    full_resp = full_resp.replace(BLOCKED_MESSAGE, "").strip()
    full_resp = full_resp.replace("https://storage.googleapis.com", "https://cdn.snapzion.com")
    final_links.extend(re.findall(r"(https://cdn\.snapzion\.com[^\s\)]+)", full_resp))
    clean = strip_model_prefix(full_resp, prefix)
    completion_tokens = calculate_tokens(clean, request.model)

    upload_replaced_urls_to_r2(final_links, alt_text=get_last_user_prompt(request.messages))

    return {
        "id": request_id,
        "object": "chat.completion",
        "created": int(datetime.now().timestamp()),
        "model": request.model,
        "system_fingerprint": system_fp,
        "choices": [{"index": 0, "message": {"role": "assistant", "content": clean}, "finish_reason": "stop"}],
        "usage": {"prompt_tokens": prompt_tokens, "completion_tokens": completion_tokens, "total_tokens": prompt_tokens + completion_tokens},
    }