Main.py

import requests, json, os, time, base64, asyncio, io
from fastapi import FastAPI, Request, HTTPException
from collections import defaultdict
from datetime import datetime, timezone
from typing import Optional
import threading
from PIL import Image as PILImage

app = FastAPI()

# ──────────────────────────────────────────────────────────────────────────────
# GitHub Setup (telemetry)
# ──────────────────────────────────────────────────────────────────────────────
GITHUB_TOKEN   = os.environ["GITHUB_TOKEN"]
REPO           = "pythonIsFast/Text2Pen-Backend"
FILE_PATH      = "telemetry.json"
API_URL_GITHUB = f"https://api.github.com/repos/{REPO}/contents/{FILE_PATH}"
GH_HEADERS     = {"Authorization": f"token {GITHUB_TOKEN}"}

# ──────────────────────────────────────────────────────────────────────────────
# Groq API
# ──────────────────────────────────────────────────────────────────────────────
GROQ_API_KEY  = os.environ["GROQ_API_KEY"]
GROQ_URL      = "https://api.groq.com/openai/v1/chat/completions"

MODEL_LARGE  = "llama-3.3-70b-versatile"           # text
MODEL_VISION = "meta-llama/llama-4-scout-17b-16e-instruct"  # vision (images)

# ──────────────────────────────────────────────────────────────────────────────
# Rate-limit state (per-IP, resets at midnight UTC)
# ──────────────────────────────────────────────────────────────────────────────
def _today_key():
    return datetime.now(timezone.utc).strftime("%Y-%m-%d")

usage: dict[str, dict] = {}
_usage_lock = threading.Lock()
telemetry_count: dict[str, list] = defaultdict(list)

def _get_real_ip(request: Request) -> str:
    forwarded_for = request.headers.get("X-Forwarded-For")
    if forwarded_for:
        return forwarded_for.split(",")[0].strip()
    return request.client.host

def _get_usage(ip: str) -> dict:
    today = _today_key()
    with _usage_lock:
        if ip not in usage or usage[ip]["date"] != today:
            usage[ip] = {"date": today, "chat": 0, "text": 0}
        return usage[ip]

def _increment(ip: str, field: str) -> int:
    today = _today_key()
    with _usage_lock:
        if ip not in usage or usage[ip]["date"] != today:
            usage[ip] = {"date": today, "chat": 0, "text": 0}
        usage[ip][field] += 1
        return usage[ip][field]

def _decrement(ip: str, field: str):
    with _usage_lock:
        if ip in usage and usage[ip].get(field, 0) > 0:
            usage[ip][field] -= 1

CHAT_LIMIT = 10
TEXT_LIMIT  = 1000

# ──────────────────────────────────────────────────────────────────────────────
# Queue infrastructure
# ──────────────────────────────────────────────────────────────────────────────
import queue as queue_module

large_queue:  queue_module.Queue = queue_module.Queue()
vision_queue: queue_module.Queue = queue_module.Queue()

def _make_job(payload: dict, model_type: str) -> dict:
    return {
        "payload":    payload,
        "model_type": model_type,
        "event":      threading.Event(),
        "result":     None,
        "error":      None,
    }

# ──────────────────────────────────────────────────────────────────────────────
# Global Groq rate limiter
# Groq free tier: ~30 req/min → 1 req per 2s to be safe
# ──────────────────────────────────────────────────────────────────────────────
class RateLimitError(Exception):
    pass

_groq_lock     = threading.Lock()
_last_call_time = 0.0
_MIN_INTERVAL   = 2.0  # seconds between Groq calls

def _call_groq(model: str, messages: list, max_tokens: int = 1000) -> str:
    """Single Groq API call with global rate limiting."""
    global _last_call_time
    with _groq_lock:
        gap = _MIN_INTERVAL - (time.time() - _last_call_time)
        if gap > 0:
            time.sleep(gap)
        _last_call_time = time.time()

        resp = requests.post(
            GROQ_URL,
            headers={"Authorization": f"Bearer {GROQ_API_KEY}",
                     "Content-Type": "application/json"},
            json={"model": model, "messages": messages, "max_tokens": max_tokens},
            timeout=60
        )

    if resp.status_code == 429:
        raise RateLimitError()
    if resp.status_code != 200:
        raise Exception(f"Groq error {resp.status_code}: {resp.text[:300]}")
    return resp.json()["choices"][0]["message"]["content"].strip()


def _call_with_backoff(model: str, messages: list, max_tokens: int = 1000) -> str:
    """Retry with exponential backoff on 429."""
    backoff = 5.0
    for attempt in range(20):
        try:
            return _call_groq(model, messages, max_tokens)
        except RateLimitError:
            if attempt == 19:
                raise
            time.sleep(backoff)
            backoff = min(backoff * 1.5, 30.0)
    raise RateLimitError()

# ──────────────────────────────────────────────────────────────────────────────
# Worker threads
# ──────────────────────────────────────────────────────────────────────────────
def _worker(q: queue_module.Queue, model: str):
    while True:
        job = q.get()
        try:
            job["result"] = _call_with_backoff(
                model,
                job["payload"]["messages"],
                job["payload"].get("max_tokens", 1000)
            )
        except RateLimitError:
            job["error"] = "rate_limit_exceeded_after_retries"
        except Exception as e:
            job["error"] = str(e)
        job["event"].set()
        q.task_done()

threading.Thread(target=_worker, args=(large_queue,  MODEL_LARGE),
                 daemon=True, name="large-worker").start()
threading.Thread(target=_worker, args=(vision_queue, MODEL_VISION),
                 daemon=True, name="vision-worker").start()


def _enqueue_and_wait(q: queue_module.Queue, payload: dict, model_type: str) -> str:
    job = _make_job(payload, model_type)
    q.put(job)
    finished = job["event"].wait(timeout=300)
    if not finished:
        raise Exception("Request timed out (5 min)")
    if job["error"]:
        if job["error"] == "rate_limit_exceeded_after_retries":
            raise RateLimitError()
        raise Exception(job["error"])
    return job["result"]

def _queue_position(q: queue_module.Queue) -> int:
    return q.qsize()

# ──────────────────────────────────────────────────────────────────────────────
# Image helpers
# ──────────────────────────────────────────────────────────────────────────────
def _get_image_mime(b64_str: str) -> str:
    try:
        raw = base64.b64decode(b64_str[:64])
        if raw[:8] == b'\x89PNG\r\n\x1a\n':          return "image/png"
        if raw[:3] == b'\xff\xd8\xff':                return "image/jpeg"
        if raw[:6] in (b'GIF87a', b'GIF89a'):         return "image/gif"
        if raw[:4] == b'RIFF' and raw[8:12] == b'WEBP': return "image/webp"
        if raw[:2] == b'BM':                          return "image/bmp"
    except Exception:
        pass
    return "image/jpeg"

def _resize_if_needed(image_b64: str) -> str:
    """Always compress to max 512x512 / JPEG Q60 to minimise token usage on Groq."""
    try:
        raw = base64.b64decode(image_b64)
        img = PILImage.open(io.BytesIO(raw))
        img.thumbnail((512, 512))
        if img.mode != "RGB":
            img = img.convert("RGB")
        buf = io.BytesIO()
        img.save(buf, format="JPEG", quality=60)
        return base64.b64encode(buf.getvalue()).decode()
    except Exception:
        return image_b64

def _build_image_part(b64_str: str) -> dict:
    mime = _get_image_mime(b64_str)
    return {"type": "image_url", "image_url": {"url": f"data:{mime};base64,{b64_str}"}}

# ──────────────────────────────────────────────────────────────────────────────
# Message builder
# ──────────────────────────────────────────────────────────────────────────────
SYSTEM_PROMPT = (
    "You are Text2Pen AI, a helpful assistant integrated into the Text2Pen desktop app. "
    "Text2Pen lets users learn their own handwriting style and reproduce typed text as "
    "handwriting in Microsoft OneNote via mouse simulation. "
    "Help users with text composition, grammar, style, and image analysis. "
    "When you suggest replacement text, make it clear so the user can click "
    "'Insert into text field' to adopt it. "
    "Keep responses concise and friendly. Respond in the same language the user writes in."
)

def _build_messages(history: list, user_text: str, image_b64: Optional[str]) -> list:
    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
    for entry in history:
        role    = entry.get("role", "user")
        content = entry.get("content", "")
        img     = entry.get("image_b64")
        if img:
            messages.append({"role": role, "content": [
                {"type": "text", "text": content},
                _build_image_part(img)
            ]})
        else:
            messages.append({"role": role, "content": content})

    if image_b64:
        messages.append({"role": "user", "content": [
            {"type": "text", "text": user_text.strip() or "What do you see in this image?"},
            _build_image_part(image_b64)
        ]})
    else:
        messages.append({"role": "user", "content": user_text})
    return messages

# ──────────────────────────────────────────────────────────────────────────────
# Endpoints
# ──────────────────────────────────────────────────────────────────────────────

@app.get("/ping")
def ping():
    return {"message": "All systems are up and running!"}


@app.post("/telemetry")
async def receive_telemetry(request: Request):
    ip  = _get_real_ip(request)
    now = time.time()

    telemetry_count[ip] = [t for t in telemetry_count[ip] if now - t < 3600]
    if len(telemetry_count[ip]) >= 100:
        raise HTTPException(status_code=429, detail="Rate limit exceeded")
    telemetry_count[ip].append(now)

    payload = await request.json()
    event = {
        "event":     payload.get("event"),
        "os":        payload.get("os"),
        "timestamp": payload.get("timestamp"),
        "extra":     payload.get("extra", {})
    }

    resp = requests.get(API_URL_GITHUB, headers=GH_HEADERS)
    if resp.status_code == 200:
        content = json.loads(base64.b64decode(resp.json()["content"]).decode())
        sha = resp.json()["sha"]
        content.append(event)
    else:
        content, sha = [event], None

    merged = {}
    for e in content:
        key = (e["event"], e["os"], tuple(sorted(e["extra"].items())))
        if key not in merged or e["timestamp"] > merged[key]["timestamp"]:
            merged[key] = e
    content = list(merged.values())

    data = {
        "message": f"Add telemetry event {int(time.time())}",
        "content": base64.b64encode(json.dumps(content, indent=2).encode()).decode()
    }
    if sha:
        data["sha"] = sha
    requests.put(API_URL_GITHUB, headers=GH_HEADERS, json=data)
    return {"status": "ok"}


@app.post("/ai")
async def ai_action(request: Request):
    ip = _get_real_ip(request)
    u  = _get_usage(ip)

    if u["text"] >= TEXT_LIMIT:
        raise HTTPException(status_code=429,
            detail=f"Daily text-action limit of {TEXT_LIMIT} reached")

    payload = await request.json()
    action  = payload.get("action")
    text    = payload.get("text", "")
    extra   = payload.get("extra", {})

    prompts = {
        "correct":     f"Correct spelling and grammar. Return ONLY the corrected text:\n\n{text}",
        "shorten":     f"Shorten this text significantly. Return ONLY the shortened text:\n\n{text}",
        "tone_formal": f"Rewrite formally. Return ONLY the rewritten text:\n\n{text}",
        "tone_casual": f"Rewrite casually/friendly. Return ONLY the rewritten text:\n\n{text}",
        "generate":    f"Write a short handwritten-style text about: {text}. Return ONLY the text.",
        "replace_unknown": (
            f"Rewrite this text so it only uses characters from this list: {extra.get('known_chars','')}. "
            f"Replace special characters with simpler alternatives (em-dash → -, semicolon → ,). "
            f"Replace any character NOT in the list with the closest equivalent that IS in the list. "
            f"Do not remove words, only substitute characters. Return ONLY the rewritten text:\n\n{text}"
        ),
    }

    prompt = prompts.get(action)
    if not prompt:
        raise HTTPException(status_code=400, detail="Unknown action")

    _increment(ip, "text")
    messages = [{"role": "user", "content": prompt}]

    def run_sync():
        return _enqueue_and_wait(large_queue, {"messages": messages}, "large")

    try:
        result = await asyncio.get_running_loop().run_in_executor(None, run_sync)
    except RateLimitError:
        _decrement(ip, "text")
        raise HTTPException(status_code=429,
            detail="Rate limit – please try again in a moment")
    except Exception as e:
        _decrement(ip, "text")
        raise HTTPException(status_code=500, detail=str(e))

    return {"result": result}


@app.post("/ai/chat")
async def ai_chat(request: Request):
    ip = _get_real_ip(request)
    u  = _get_usage(ip)

    if u["chat"] >= CHAT_LIMIT:
        raise HTTPException(status_code=429,
            detail=f"Daily chat limit of {CHAT_LIMIT} reached ({u['chat']}/{CHAT_LIMIT})")

    payload   = await request.json()
    history   = payload.get("history", [])
    user_text = payload.get("message", "")
    image_b64 = payload.get("image_b64") or None

    if not user_text.strip() and not image_b64:
        raise HTTPException(status_code=400, detail="Empty message")

    if image_b64:
        image_b64 = _resize_if_needed(image_b64)

    use_vision = image_b64 is not None
    q          = vision_queue if use_vision else large_queue
    messages   = _build_messages(history, user_text, image_b64)
    pos        = _queue_position(q) + 1

    _increment(ip, "chat")

    def run_sync():
        return _enqueue_and_wait(q, {"messages": messages}, "vision" if use_vision else "large")

    try:
        result = await asyncio.get_running_loop().run_in_executor(None, run_sync)
    except RateLimitError:
        _decrement(ip, "chat")
        raise HTTPException(status_code=429,
            detail="Rate limit – please wait a moment and try again")
    except Exception as e:
        _decrement(ip, "chat")
        raise HTTPException(status_code=500, detail=str(e))

    u = _get_usage(ip)
    return {
        "result":         result,
        "model":          "llama-4-vision" if use_vision else "llama-3.3-70b",
        "queue_position": pos,
        "chat_used":      u["chat"],
        "chat_limit":     CHAT_LIMIT,
    }


@app.get("/ai/limits")
async def get_limits(request: Request):
    ip = _get_real_ip(request)
    u  = _get_usage(ip)
    return {
        "chat_used":          u["chat"],
        "chat_limit":         CHAT_LIMIT,
        "text_used":          u["text"],
        "text_limit":         TEXT_LIMIT,
        "resets_at":          "midnight UTC",
        "large_queue_depth":  large_queue.qsize(),
        "vision_queue_depth": vision_queue.qsize(),
    }