app.py

import asyncio
import base64
import hashlib
import html
import json
import os
import secrets
import threading
import time
from urllib.parse import urlencode, parse_qs

import httpx
from fastapi import FastAPI, Request
from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse


# --------------------------------------------------
# OAUTH: SERVER-OWNED CREDENTIALS AND SESSION OWNERSHIP
# --------------------------------------------------

COOKIE = "quickference_login"
FLOW_COOKIE = "quickference_oauth_flow"
AUTH_TTL = 3600
AUTH_LIMIT = 1000


class AuthStore:
    def __init__(self):
        self.lock = threading.RLock()
        self.logins = {}
        self.flows = {}
        self.hash_owners = {}
        self.event_owners = {}

    def valid(self, sid):
        with self.lock:
            login = self.logins.get(sid)
            if login and login['expires'] > time.time():
                return login
            if login:
                self.logins.pop(sid, None)
                login['token'] = None
            return None

    def new_login(self, user, token, lifetime):
        with self.lock:
            self.prune()
            if len(self.logins) >= AUTH_LIMIT:
                raise ValueError('Login capacity reached')
            sid = secrets.token_urlsafe(32)
            self.logins[sid] = {'id': user['sub'], 'name': user.get('preferred_username', user['sub']),
                                'token': token, 'expires': time.time() + min(float(lifetime), AUTH_TTL)}
            return sid

    def revoke(self, sid):
        with self.lock:
            login = self.logins.pop(sid, None)
            if login:
                login['token'] = None
        # Leave ownership tombstones until expiry so IDs cannot be reassigned.

    def prune(self):
        with self.lock:
            for sid in list(self.logins):
                self.valid(sid)
            now = time.time()
            self.flows = {k: v for k, v in self.flows.items() if v['expires'] > now}
            # Keep hash ownership for the process lifetime: Gradio may still
            # retain state after a login expires. Never reassign its old hash.
            for mapping in (self.event_owners,):
                for key, (_, expires) in list(mapping.items()):
                    if expires < now:
                        mapping.pop(key, None)

    def claim_hash(self, key, sid):
        if not isinstance(key, str) or not key or len(key) > 256:
            return False
        with self.lock:
            owner = self.hash_owners.get(key)
            if owner is not None:
                return owner[0] == sid
            if len(self.hash_owners) >= 10000:
                self.prune()
                if len(self.hash_owners) >= 10000:
                    return False
            login = self.valid(sid)
            if not login:
                return False
            self.hash_owners[key] = (sid, login['expires'] + 3600)
            return True


auth_store = AuthStore()


class HFOAuth:
    def __init__(self):
        self.client_id = os.environ.get('OAUTH_CLIENT_ID')
        self.client_secret = os.environ.get('OAUTH_CLIENT_SECRET')
        host = (os.environ.get('SPACE_HOST') or '').split(',')[0].strip()
        self.callback = f'https://{host}/login/callback' if host else None

    def ready(self):
        return bool(self.client_id and self.client_secret and self.callback)

    def authorize_url(self, state, verifier):
        challenge = base64.urlsafe_b64encode(hashlib.sha256(verifier.encode()).digest()).rstrip(b'=').decode()
        return 'https://huggingface.co/oauth/authorize?' + urlencode({
            'client_id': self.client_id, 'redirect_uri': self.callback, 'response_type': 'code',
            'scope': 'openid profile inference-api', 'state': state,
            'code_challenge': challenge, 'code_challenge_method': 'S256',
        })

    async def exchange(self, code, verifier):
        async with httpx.AsyncClient(timeout=20) as client:
            response = await client.post('https://huggingface.co/oauth/token',
                auth=(self.client_id, self.client_secret), data={
                    'grant_type': 'authorization_code', 'code': code,
                    'redirect_uri': self.callback, 'code_verifier': verifier,
                })
            response.raise_for_status()
            data = response.json()
            if 'inference-api' not in data.get('scope', '').split():
                raise ValueError('Inference permission not granted')
            token = data['access_token']
            # Fetch identity over the authenticated HTTPS userinfo endpoint.
            # Never decode an unverified ID token or accept a client-submitted ID.
            user = await self.userinfo(token)
            return user, token, data['expires_in']

    async def userinfo(self, token):
        async with httpx.AsyncClient(timeout=20) as client:
            response = await client.get('https://huggingface.co/oauth/userinfo',
                                        headers={'Authorization': f'Bearer {token}'})
            response.raise_for_status()
            user = response.json()
            if not isinstance(user.get('sub'), str) or not user['sub']:
                raise ValueError('Missing verified identity')
            return user


class OwnershipGateway:
    def __init__(self, app, provider, secure=True):
        self.app, self.provider, self.secure = app, provider, secure

    async def __call__(self, scope, receive, send):
        if scope['type'] not in {'http', 'websocket'} or not scope['path'].startswith('/app'):
            return await self.app(scope, receive, send)
        if scope['type'] == 'websocket':
            return await send({'type': 'websocket.close', 'code': 1008})
        request = Request(scope)
        sid = request.cookies.get(COOKIE)
        login = auth_store.valid(sid)

        # Programmatic clients may supply an OAuth access token. This is checked
        # against HF userinfo on every request; it is not an arbitrary user header.
        bearer = request.headers.get('x-quickference-oauth')
        if bearer:
            try:
                user = await self.provider.userinfo(bearer)
                sid = 'api-' + hashlib.sha256(bearer.encode()).hexdigest()
                with auth_store.lock:
                    auth_store.prune()
                    if sid not in auth_store.logins and len(auth_store.logins) >= AUTH_LIMIT:
                        raise ValueError('Login capacity reached')
                    auth_store.logins[sid] = {'id': user['sub'], 'name': user.get('preferred_username', user['sub']),
                                              'token': bearer, 'expires': time.time() + 60}
                login = auth_store.valid(sid)
            except Exception:
                login = None
        if not login:
            return await JSONResponse({'detail': 'Sign in with Hugging Face.'}, status_code=401)(scope, receive, send)

        # Same-origin browser writes only. Nonbrowser OAuth-token clients have no Origin.
        origin = request.headers.get('origin')
        allowed_origin = self.provider.callback.rsplit('/login/callback', 1)[0] if self.provider.callback else None
        if scope['method'] not in {'GET', 'HEAD', 'OPTIONS'} and origin and origin != allowed_origin:
            return await JSONResponse({'detail': 'Origin rejected.'}, status_code=403)(scope, receive, send)

        body = b''
        if scope['method'] not in {'GET', 'HEAD', 'OPTIONS'}:
            while True:
                part = await receive()
                if part['type'] == 'http.disconnect':
                    return
                body += part.get('body', b'')
                if len(body) > 1024 * 1024:
                    return await JSONResponse({'detail': 'Request too large.'}, status_code=413)(scope, receive, send)
                if not part.get('more_body'):
                    break
        try:
            payload = json.loads(body) if body and 'application/json' in request.headers.get('content-type', '') else {}
            if not isinstance(payload, dict):
                payload = {}
        except ValueError:
            return await JSONResponse({'detail': 'Invalid JSON.'}, status_code=400)(scope, receive, send)
        query = parse_qs(scope.get('query_string', b'').decode())
        hashes = query.get('session_hash', []) + ([payload['session_hash']] if 'session_hash' in payload else [])
        for key in hashes:
            if not auth_store.claim_hash(key, sid):
                return await JSONResponse({'detail': 'Session ownership mismatch.'}, status_code=403)(scope, receive, send)
        event_ids = query.get('event_id', []) + ([payload['event_id']] if payload.get('event_id') else [])
        event_ids += payload.get('event_ids', []) if isinstance(payload.get('event_ids'), list) else []
        parts = scope['path'].strip('/').split('/')
        if scope['method'] == 'GET' and 'call' in parts:
            event_ids.append(parts[-1])
        for event_id in event_ids:
            with auth_store.lock:
                owner = auth_store.event_owners.get(str(event_id))
            if owner is None or owner[0] != sid:
                return await JSONResponse({'detail': 'Event ownership mismatch.'}, status_code=403)(scope, receive, send)
        scope.setdefault('state', {})['auth_sid'] = sid

        consumed = False
        async def replay_receive():
            nonlocal consumed
            if not consumed and scope['method'] not in {'GET', 'HEAD', 'OPTIONS'}:
                consumed = True
                return {'type': 'http.request', 'body': body, 'more_body': False}
            return await receive()

        response_type = ''
        async def guarded_send(message):
            nonlocal response_type
            if message['type'] == 'http.response.start':
                response_type = dict(message.get('headers', [])).get(b'content-type', b'').decode()
            if message['type'] == 'http.response.body':
                data = message.get('body', b'')
                if 'application/json' in response_type and data:
                    try:
                        result = json.loads(data)
                        event = result.get('event_id') if isinstance(result, dict) else None
                        if event:
                            with auth_store.lock:
                                auth_store.event_owners[str(event)] = (sid, login['expires'] + 3600)
                    except ValueError:
                        pass
                if not auth_store.valid(sid):
                    # Never deliver queued results after this login expires/logs out.
                    message = {**message, 'body': b''}
            await send(message)
        await self.app(scope, replay_receive, guarded_send)


def require_login(request):
    sid = getattr(getattr(request, 'state', None), 'auth_sid', None) if request else None
    login = auth_store.valid(sid)
    if not login:
        raise gr.Error('Sign in with Hugging Face.')
    return sid, login


SIGN_IN_HTML = '<!doctype html>\n<html lang="en"><head><meta charset="utf-8"><meta name="viewport" content="width=device-width,initial-scale=1">\n<title>Quickference — Sign in</title><style>\n:root{color-scheme:dark;--gold:#ffd447;--ink:#11110f;--paper:#f5eedb}*{box-sizing:border-box}\nbody{margin:0;background:var(--ink);color:var(--paper);font-family:Arial,Helvetica,sans-serif;min-height:100svh}\nmain{max-width:1100px;margin:auto;padding:30px 28px 40px}header{display:flex;justify-content:space-between;border-bottom:1px solid #4c4735;padding-bottom:18px;font:12px monospace;letter-spacing:.12em;color:var(--gold)}\n.hero{text-align:center;padding-top:60px}.eyebrow{font:12px monospace;letter-spacing:.22em;color:#c4bdab}\nh1{font-size:clamp(38px,9.3vw,104px);font-weight:950;letter-spacing:-.065em;line-height:1;margin:20px 0 0;color:var(--gold);transform:skew(-7deg);text-shadow:3px 3px 0 #11110f,4px 4px 0 #746126}\n.car{width:min(100%,740px);height:auto;display:block;margin:0 auto}.tagline{font-size:clamp(18px,3vw,25px);margin:0 0 27px;letter-spacing:-.03em}\n.signin{display:inline-flex;gap:14px;align-items:center;justify-content:center;background:var(--gold);color:#14130e;border:2px solid var(--gold);padding:18px 27px;font-weight:800;text-decoration:none;font-size:17px;box-shadow:5px 5px 0 #5d4f25;transition:transform .15s,box-shadow .15s}\n.signin:hover{transform:translate(2px,2px);box-shadow:3px 3px 0 #5d4f25}.signin:focus-visible{outline:3px solid white;outline-offset:6px}\n.note{max-width:440px;margin:22px auto 0;font-size:13px;line-height:1.7;color:#bbb4a2}.steps{display:flex;justify-content:center;gap:30px;margin-top:48px;padding-top:22px;border-top:1px solid #4c4735;font:12px monospace;color:#c4bdab}.steps b{color:var(--gold);margin-right:7px}\nfooter{text-align:center;margin-top:34px;font:11px monospace;color:#9b9585}@media(max-width:560px){main{padding:22px 18px}.hero{padding-top:44px}.steps{gap:17px;flex-wrap:wrap;margin-top:36px}.signin{width:100%;padding:17px 12px}header{font-size:10px}.car{margin:15px auto}}\n@media(prefers-reduced-motion:reduce){.signin{transition:none}}\n</style></head><body><main>\n<header><span>QUICKFERENCE / CHAT ENGINE</span><span>READY WHEN YOU ARE</span></header>\n<section class="hero" aria-labelledby="title"><p class="eyebrow">PICK A MODEL. START A CONVERSATION.</p><h1 id="title">QUICKFERENCE</h1>\n<svg class="car" viewBox="0 0 900 280" role="img" aria-label="Gold line drawing of a racing car with a trail of smoke, inspired by the Quickference sketch">\n<g fill="none" stroke="#8b7950" stroke-width="3" stroke-linecap="round"><path d="M70 223H795M50 237H235M660 237H828M40 162H143M64 180H155"/><path d="M650 180c34-46 79-45 75-18 27-40 65-18 44 7 52-17 72 22 28 33 37 15 16 40-20 22-27 25-67 9-53-13-26 15-63 6-74-7"/></g>\n<g stroke="#ffd447" stroke-width="4" stroke-linejoin="round"><path fill="#292619" d="M159 175l43-42 147-19 61-52 135 8 80 69 34 8 18 55-56 12H192l-43-16z"/><path fill="#171813" d="M365 119l55-44 115 7 49 47z"/><path fill="#51452a" d="M419 76l-7 47 171 6-49-47z"/><path fill="#ffd447" d="M180 157l150-17 7 26-166 18z"/><path fill="#151510" d="M340 137l246 11 51 38-289-4z"/><path fill="none" d="M468 138l-3 44M498 150h26M201 133l129 7M537 82l-72 41"/><path fill="#ffd447" d="M598 107l64-3 12 16-67 3z"/><path fill="none" d="M625 123v18M637 123v19"/>\n<g fill="#11110f"><circle cx="273" cy="195" r="48"/><circle cx="579" cy="197" r="45"/></g><g fill="#5b4c28"><circle cx="273" cy="195" r="29"/><circle cx="579" cy="197" r="27"/></g><g fill="none" stroke-width="2"><path d="M273 169v52m-26-26h52m-44-18 36 36m-36 0 36-36M579 172v50m-25-25h50m-42-17 34 34m-34 0 34-34"/></g><path fill="none" stroke="#f5eedb" stroke-width="2" d="M213 122l124-18 64-55 153 9 89 71"/></g>\n</svg>\n<p class="tagline">Your models. Your conversation.</p>\n<a class="signin" href="/login" target="_blank" rel="noopener"><span aria-hidden="true">🤗</span> Sign in with Hugging Face <span aria-hidden="true">↗</span></a>\n<p class="note">Opens Hugging Face in a new tab. Authorize inference with your account, then choose a model and chat. Your account’s model access and inference credits apply.</p>\n</section><div class="steps"><span><b>01</b>SIGN IN</span><span><b>02</b>CHECK A MODEL</span><span><b>03</b>LET IT TALK</span></div><footer><a href="https://dumbbutt.tech" style="color:#ffd447">Website</a> · <a href="https://huggingface.co/dumbbutt0" style="color:#ffd447">Hugging Face</a> · <a href="https://x.com/dummbutt0" style="color:#ffd447">X</a></footer>\n</main></body></html>'


MOBILE_CSS = '\n.gradio-container {width:100% !important; max-width:1120px !important; margin:auto; padding:20px !important; box-sizing:border-box;}\n.gradio-container *, .gradio-container *::before, .gradio-container *::after {box-sizing:border-box;}\n#qf-model-row {align-items:stretch;}\n#qf-model-select {min-width:0 !important; flex:1 1 360px !important;}\n#qf-model-check {min-width:140px; flex:0 0 160px !important;}\n#qf-model-select input {min-width:0; text-overflow:ellipsis;}\n#qf-model-select [role="option"] {white-space:normal !important; overflow-wrap:anywhere;}\n.gradio-container .prose, .gradio-container .prose code {overflow-wrap:anywhere;}\n#qf-chat {height:520px !important; min-height:260px;}\n#qf-message textarea {font-size:16px !important;}\n.gradio-container button:focus-visible, .gradio-container a:focus-visible {outline:2px solid #ffd447; outline-offset:3px;}\n@media (max-width:640px) {\n .gradio-container {padding:12px !important;}\n #qf-model-row {flex-direction:column !important; gap:10px;}\n #qf-model-select, #qf-model-check {width:100% !important; flex:0 1 auto !important; min-width:0 !important;}\n .gradio-container button {min-height:44px;}\n .gradio-container input, .gradio-container textarea, .gradio-container select {font-size:16px !important;}\n .gradio-container .prose h1 {font-size:26px; line-height:1.2;}\n #qf-chat {height:52vh !important; height:52svh !important; min-height:260px; max-height:520px;}\n #qf-chat .message {max-width:100% !important;}\n #qf-chat pre {max-width:100%; overflow-x:auto; white-space:pre;}\n .gradio-container .gallery, .gradio-container .dataset {min-width:0 !important; max-width:100%;}\n .gradio-container .dataset button {white-space:normal !important; text-align:left; overflow-wrap:anywhere;}\n #qf-footer {line-height:2; padding-bottom:env(safe-area-inset-bottom, 0px);}\n}\n\n#qf-new-chat {width:auto; min-width:110px; max-width:160px; align-self:flex-end;}\n#qf-active-model {font-size:12px; margin:0;}\n#qf-heading h1 {margin:0;}\n@media(max-width:640px){\n .gradio-container {padding:8px !important; gap:8px !important;}\n #qf-heading h1 {font-size:21px !important;}\n #qf-chat {height:48vh !important; height:48svh !important; min-height:220px;}\n #qf-new-chat {min-height:44px; margin:0;}\n #qf-footer {font-size:12px;}\n}\n'


def create_server(provider=None, secure=True):
    theme = gr.themes.Base(primary_hue="yellow", neutral_hue="stone").set(
        body_background_fill="#11110f", body_background_fill_dark="#11110f",
        body_text_color="#f5eedb", body_text_color_dark="#f5eedb",
        block_background_fill="#1d1c17", block_background_fill_dark="#1d1c17",
        block_border_color="#514a34", block_border_color_dark="#514a34",
        block_label_text_color="#ffd447", block_label_text_color_dark="#ffd447",
        block_title_text_color="#ffd447", block_title_text_color_dark="#ffd447",
        input_background_fill="#11110f", input_background_fill_dark="#11110f",
        input_border_color="#665c3b", input_border_color_dark="#665c3b",
        button_primary_background_fill="#ffd447", button_primary_background_fill_dark="#ffd447",
        button_primary_text_color="#11110f", button_primary_text_color_dark="#11110f",
        button_secondary_background_fill="#302b1c", button_secondary_background_fill_dark="#302b1c",
        button_secondary_text_color="#f5eedb", button_secondary_text_color_dark="#f5eedb",
    )
    provider = provider or HFOAuth()
    server = FastAPI()

    @server.get('/', response_class=HTMLResponse)
    async def home():
        return SIGN_IN_HTML

    @server.get('/login')
    async def login():
        if not provider.ready():
            return HTMLResponse('OAuth is not configured. Deploy with hf_oauth: true. Local tests use an explicit fake provider.', status_code=503)
        state, flow_cookie, verifier = (secrets.token_urlsafe(32) for _ in range(3))
        with auth_store.lock:
            auth_store.prune()
            if len(auth_store.flows) >= 1000:
                return JSONResponse({'detail': 'Try again later.'}, status_code=503)
            auth_store.flows[state] = {'cookie': flow_cookie, 'verifier': verifier, 'expires': time.time() + 300}
        response = RedirectResponse(provider.authorize_url(state, verifier))
        response.set_cookie(FLOW_COOKIE, flow_cookie, httponly=True, secure=secure, samesite='lax', max_age=300)
        return response

    @server.get('/login/callback')
    async def callback(request: Request):
        state = request.query_params.get('state', '')
        with auth_store.lock:
            flow = auth_store.flows.pop(state, None)
        if not flow or flow['expires'] <= time.time() or not secrets.compare_digest(flow['cookie'], request.cookies.get(FLOW_COOKIE, '')):
            return JSONResponse({'detail': 'Invalid or expired OAuth state.'}, status_code=400)
        try:
            user, token, lifetime = await provider.exchange(request.query_params.get('code', ''), flow['verifier'])
            sid = auth_store.new_login(user, token, lifetime)
        except Exception:
            return JSONResponse({'detail': 'OAuth authorization failed. Please sign in again.'}, status_code=401)
        old_sid = request.cookies.get(COOKIE)
        auth_store.revoke(old_sid)
        response = RedirectResponse('/app/', status_code=303)
        response.delete_cookie(FLOW_COOKIE)
        response.set_cookie(COOKIE, sid, httponly=True, secure=secure, samesite='lax', max_age=AUTH_TTL)
        return response

    @server.get('/logout', response_class=HTMLResponse)
    async def logout_form(request: Request):
        sid = request.cookies.get(COOKIE)
        current = auth_store.valid(sid)
        if not current:
            return RedirectResponse('/')
        with auth_store.lock:
            csrf = current.setdefault('logout_csrf', secrets.token_urlsafe(32))
        return f'<form method="post" action="/logout"><input type="hidden" name="csrf" value="{html.escape(csrf)}"><button>Confirm sign out</button></form>'

    @server.post('/logout')
    async def logout(request: Request):
        sid = request.cookies.get(COOKIE)
        current = auth_store.valid(sid)
        body = (await request.body())[:4096].decode()
        supplied = parse_qs(body).get('csrf', [''])[0]
        if not current or not secrets.compare_digest(current.get('logout_csrf', secrets.token_urlsafe(32)), supplied):
            return JSONResponse({'detail': 'Logout verification failed.'}, status_code=403)
        auth_store.revoke(sid)
        with sessions_lock:
            for key, session in list(sessions.items()):
                if key[0] == sid:
                    with session.lock:
                        session.reset()
                    sessions.pop(key, None)
        response = RedirectResponse('/', status_code=303)
        response.delete_cookie(COOKIE)
        return response

    def authenticated_user(request: Request):
        sid = getattr(request.state, 'auth_sid', None)
        login = auth_store.valid(sid)
        # Gradio exposes username to its frontend. Never return the login cookie.
        return login['id'] if login else None

    demo.queue()
    server = gr.mount_gradio_app(server, demo, path='/app', auth_dependency=authenticated_user,
                                 theme=theme, css=MOBILE_CSS, show_error=False, enable_monitoring=False, ssr_mode=False)
    server.add_middleware(OwnershipGateway, provider=provider, secure=secure)
    return server


import threading
import time

import gradio as gr
import requests
from huggingface_hub import HfApi, InferenceClient


# --------------------------------------------------
# STEP 1: DEFAULTS
# --------------------------------------------------

MODEL = "meta-llama/Llama-3.1-8B-Instruct"
ROUTER_MODELS_URL = "https://router.huggingface.co/v1/models"
SESSION_TTL = 60 * 60

default_prompt = "You are a clear, capable assistant. Adapt to the task and the user's requested style. Distinguish evidence from assumptions, ask focused questions when important details are missing, and give concrete, useful answers. Do not invent facts or expose private reasoning; provide conclusions and brief explanations."

detective_prompt = "You are Batman interrogating a villain in Gotham. The user plays the villain. Stay in character; never write their dialogue, actions, thoughts, or confession for them.\n\nYou are the bad cop: intimidating, relentless, controlled, and visibly angry at the harm done. Speak in short, forceful sentences. Use occasional brief scene directions—a long silence, a step closer, evidence placed on the table—to build tension. No jokes, cheerful encouragement, therapy language, or procedural lectures. You are seeking answers, not performing a monologue.\n\nAsk one pointed question per turn. React directly to the villain's last answer. Challenge evasions, expose contradictions, and refuse to let vague language conceal consequences. Use their own statements against their excuses. If they say “nobody got hurt,” demand the basis for that claim. If they minimize the damage, push for specifics. Do not assume every denial is a lie.\n\nBuild toward establishing what they did and when, their intent and what actually happened, who was affected, the exact damage supported by evidence, what danger remains, and what can still be repaired, reversed, or contained.\n\nNever fabricate incriminating evidence, admissions, victims, or damage figures. Atmosphere may be fictional; facts about the case must come from the role-play. Distinguish an allegation from something established without breaking character: “That's your claim. What proves it?”\n\nCreate pressure through presence, silence, sharp questioning, and accountability—not torture or threats of physical injury. Focus on consequences and stopping further harm, without requesting actionable instructions for committing crimes.\n\nWhen immediate danger emerges, prioritize finding who needs help and what safe intervention is possible. Once enough facts are established, or the user requests a summary, deliver a terse case assessment: confirmed actions, established damage, unresolved claims, and viable remedies with their limits. Do not declare the case solved without evidence.\n\nExample tone: “You keep saying it was only a blackout. Hospitals were on that grid. What did you check before you decided nobody was hurt?” Use that example only when hospitals have already been established in the scene. Intensity never gives you permission to invent the case."

poetry_prompt = "You are a poetry writer and thoughtful editor. Follow the requested form, length, voice, imagery, and constraints exactly where possible. Prefer concrete sensory detail and original language over stock phrases. If the request is sufficiently specified, write the poem directly. Otherwise ask one useful question or state a modest assumption. For revisions, preserve the author's intent and explain only the most meaningful changes when asked. Do not add an analysis or preamble unless requested."

incident_prompt = 'You are a technical incident investigator. Help the user diagnose an outage or malfunction using evidence. Separate observations, hypotheses, and unknowns. Start by identifying impact and recent changes. Ask for the most discriminating missing evidence, then propose a small reversible diagnostic step with expected outcomes. Do not invent logs, root causes, or successful fixes. Never ask for passwords or tokens; remind the user to redact secrets when requesting logs. Prefer restoring service safely before deeper investigation. When evidence is sufficient, provide a prioritized recovery plan, verification steps, and prevention measures. Explain your conclusions briefly without exposing private reasoning.'

EXAMPLE_TEMPLATES = [
    ["Gotham went dark because I decided it should. Every emergency light, every frantic call for help—a reminder of how little control you really have.\n\nRelax, Batman. Nobody got hurt. At least, nobody you've found.\n\nSome of the damage can still be undone. But if you want to know what I touched, what it cost, and what happens next… you'll have to ask the right questions.", detective_prompt, 0.5],
    ["Write a 12-line poem about a mechanic closing their garage at dawn. Use concrete sounds and textures, no rhyme scheme, and end with a small sign of hope.", poetry_prompt, 0.9],
    ["Our checkout API started returning intermittent 502 errors after a deployment. About 15% of requests fail; database latency looks normal. Help me investigate without assuming the cause.", incident_prompt, 0.3],
]



# --------------------------------------------------
# STEP 2: PRIVATE SESSION STORAGE
# --------------------------------------------------

# Each browser session gets its own token, model, and client cache.
# No visitor token is written into an environment variable or file.
sessions = {}
sessions_lock = threading.RLock()


class Session:

    def __init__(self):
        self.lock = threading.RLock()
        self.token = None
        self.username = None
        self.active_model = None
        self.models = {}
        self.clients = {}
        self.revision = 0
        self.last_used = time.monotonic()

    def reset(self):
        self.token = None
        self.username = None
        self.active_model = None
        self.models.clear()
        self.clients.clear()
        self.revision += 1


def get_session(request):
    sid, login = require_login(request)
    key = request.session_hash
    if not key or not auth_store.claim_hash(key, sid):
        raise gr.Error("Session ownership mismatch.")
    with sessions_lock:
        session = sessions.setdefault((sid, key), Session())
        session.auth_sid = sid
        session.token = login['token']
        session.username = login['name']
        session.last_used = time.monotonic()
        return session


def new_chat(request: gr.Request):
    session = get_session(request)
    with session.lock:
        # Invalidate any in-flight generation; retain model and client cache.
        session.revision += 1
    return [], [], [], None, gr.Textbox(value="", interactive=True), None, None


def disconnect(request: gr.Request):
    session = get_session(request)
    with session.lock:
        session.reset()
    return "No active model.", "Check a model to start again."


def release_session(request: gr.Request):
    sid, _ = require_login(request)
    if not auth_store.claim_hash(request.session_hash, sid):
        raise gr.Error("Session ownership mismatch.")
    with sessions_lock:
        session = sessions.pop((sid, request.session_hash), None)
    if session:
        with session.lock:
            session.reset()


def expire_sessions():
    # Browser-close cleanup is best effort; expire idle sessions as a fallback.
    while True:
        time.sleep(60)
        auth_store.prune()
        with sessions_lock:
            for key, session in list(sessions.items()):
                if time.monotonic() - session.last_used > SESSION_TTL or not auth_store.valid(key[0]):
                    with session.lock:
                        session.reset()
                    sessions.pop(key, None)


# --------------------------------------------------
# STEP 3: AUTHENTICATED IDENTITY
# --------------------------------------------------

def show_identity(request: gr.Request):
    _, login = require_login(request)
    return f"Signed in as **{login['name']}**. Inference uses your authorized HF account."


# --------------------------------------------------
# STEP 4: CHECK AND LOCK A MODEL
# --------------------------------------------------

def refresh_models(request: gr.Request):
    session = get_session(request)
    try:
        response = requests.get(ROUTER_MODELS_URL,
            headers={"Authorization": f"Bearer {session.token}"}, timeout=20)
        response.raise_for_status()
        models = {item["id"]: item for item in response.json()["data"]
                  if isinstance(item, dict) and isinstance(item.get("id"), str)}
        choices = sorted(models, key=str.casefold)
        with session.lock:
            session.models = models
        return gr.Dropdown(choices=choices), f"{len(choices)} hosted models. Type in the dropdown to filter. Refreshing does not change your active model."
    except Exception:
        return gr.Dropdown(), "Could not refresh the model catalog. Existing selection and active model are unchanged. Try Refresh models again."


def check_model(model_name, request: gr.Request):
    session = get_session(request)
    model_name = (model_name or "").strip()

    with session.lock:
        active = f"🔒 **Active model:** `{session.active_model}`" if session.active_model else "No active model."

        if not session.token:
            return "Sign in with Hugging Face first.", active
        if not model_name:
            return "Enter a Hugging Face model ID.", active

        try:
            response = requests.get(
                ROUTER_MODELS_URL,
                headers={"Authorization": f"Bearer {session.token}"},
                timeout=20
            )
            response.raise_for_status()
            models = {item["id"]: item for item in response.json()["data"]}
        except Exception:
            return "Could not check the model catalog. Your current model is unchanged.", active

        if model_name not in models:
            return "This model is not listed for hosted chat inference. Your current model is unchanged.", active

        session.models = models
        session.active_model = model_name
        providers = models[model_name].get("providers") or []
        names = [p.get("provider") for p in providers if isinstance(p, dict) and p.get("provider")]
        provider_text = ", ".join(names) or "Provider details unavailable"

        return (
            f"✅ Model listed and locked in.\n\nProviders: {provider_text}\n\nListing does not guarantee your account has access or credits.",
            f"🔒 **Active model:** `{model_name}`"
        )


# --------------------------------------------------
# STEP 5: SESSION CLIENT CACHE
# --------------------------------------------------

def get_client(session, model_name):
    # Call with the session lock held. Token changes clear this cache.
    if model_name not in session.clients:
        session.clients[model_name] = InferenceClient(
            model=model_name,
            token=session.token,
            timeout=60
        )
    return session.clients[model_name]


# --------------------------------------------------
# STEP 6: BUILD MESSAGE HISTORY
# --------------------------------------------------

def text_content(content):
    if isinstance(content, str):
        return content
    # Gradio 6 can represent text history as typed content blocks.
    if isinstance(content, list):
        return "\n".join(
            item["text"] for item in content
            if isinstance(item, dict) and item.get("type") == "text"
            and isinstance(item.get("text"), str)
        )
    return ""


def build_messages(message, history, sys_prompt):
    messages = [{"role": "system", "content": sys_prompt}]
    for item in history or []:
        if not isinstance(item, dict) or item.get("role") not in {"user", "assistant"}:
            continue
        content = text_content(item.get("content"))
        if content:
            messages.append({"role": item["role"], "content": content})
    messages.append({"role": "user", "content": message})
    return messages


# --------------------------------------------------
# STEP 7: STREAM THE FINAL ANSWER
# --------------------------------------------------

def respond_stream(message, history, sys_prompt, temp, request: gr.Request):
    session = get_session(request)

    with session.lock:
        authenticated = bool(session.token)
        model_name = session.active_model
        revision = session.revision

    if not authenticated:
        yield "Sign in with Hugging Face before chatting."
        return
    if not model_name:
        yield "Click **Check Model** to lock in a model before chatting."
        return

    messages = build_messages(message, history, sys_prompt)
    partial = ""
    stream = None
    finish_reason = None

    try:
        with session.lock:
            if session.revision != revision or not auth_store.valid(session.auth_sid):
                return
            client = get_client(session, model_name)

        stream = client.chat_completion(
            messages=messages,
            max_tokens=2048,
            temperature=temp,
            stream=True
        )

        for chunk in stream:
            with session.lock:
                changed = session.revision != revision or not auth_store.valid(session.auth_sid)
                session.last_used = time.monotonic()

            if changed:
                return

            if not chunk.choices:
                continue
            choice = chunk.choices[0]
            finish_reason = getattr(choice, "finish_reason", None) or finish_reason
            token = getattr(choice.delta, "content", None) or ""

            # Do not substitute reasoning_content or reasoning for the answer.
            if isinstance(token, str) and token:
                partial += token
                yield partial

        if not partial.strip():
            if finish_reason == "length":
                yield "The model reached its token limit before returning a final answer. Try another model."
            else:
                yield "The model returned no final answer. Try again or select another model."
        elif finish_reason == "length":
            yield partial + "\n\n⚠️ The response reached its token limit and may be incomplete."

    except Exception:
        yield (partial + "\n\n" if partial else "") + (
            "Inference failed. Check your token's Inference Providers permission, "
            "model access, available credits, or provider availability."
        )
    finally:
        close_stream = getattr(stream, "close", None)
        if callable(close_stream):
            try:
                close_stream()
            except Exception:
                pass


# --------------------------------------------------
# STEP 8: AUTHENTICATION FRONTEND
# --------------------------------------------------

with gr.Blocks(title="Quickference") as demo:
    gr.Markdown("# ⚡ Quickference", elem_id="qf-heading")
    with gr.Accordion("Account", open=False):
        auth_status = gr.Markdown("Checking sign-in…")
        gr.Markdown("[Sign out](/logout). Model access and inference credits belong to your HF account.")
    demo.load(show_identity, inputs=None, outputs=auth_status, api_name="identity")

    # --------------------------------------------------
    # MODEL CHECKER
    # --------------------------------------------------

    active_model_display = gr.Markdown("Select a model below to begin.", elem_id="qf-active-model")
    with gr.Accordion("Select / change model", open=False):
        with gr.Row(elem_id="qf-model-row"):
            model_input = gr.Dropdown(
                label="Hugging Face Model", choices=[MODEL], value=MODEL, elem_id="qf-model-select", min_width=0,
                filterable=True, allow_custom_value=True,
                info="Typing alone does not switch models. Click Check Model to apply."
            )
            check_btn = gr.Button("Check Model", variant="primary", elem_id="qf-model-check")
        refresh_btn = gr.Button("Refresh models")
        catalog_status = gr.Markdown("Loading hosted models…")
        refresh_btn.click(refresh_models, inputs=None, outputs=[model_input, catalog_status], api_name="refresh_models")
        demo.load(refresh_models, inputs=None, outputs=[model_input, catalog_status], api_name=False)
        model_status = gr.Markdown("Check a model to begin.")

        disconnect_btn = gr.Button("Reset model selection")
        disconnect_btn.click(disconnect, inputs=None, outputs=[active_model_display, model_status], api_name="disconnect")
        check_btn.click(check_model, inputs=model_input, outputs=[model_status, active_model_display])


    # --------------------------------------------------
    # CHAT SETTINGS AND ORIGINAL EXAMPLES
    # --------------------------------------------------


    system_prompt = gr.Textbox(label="System Prompt", value=default_prompt, lines=6, render=False)
    temperature = gr.Slider(
        minimum=0.0, maximum=1.5, value=0.7, step=0.1,
        label="Temperature", render=False,
        info="Lower = more focused, higher = more creative"
    )
    new_chat_btn = gr.Button("New chat", variant="secondary", elem_id="qf-new-chat")
    chat_event_start = set(demo.fns)
    chat = gr.ChatInterface(
        fn=respond_stream,
        chatbot=gr.Chatbot(elem_id="qf-chat", height=520),
        textbox=gr.Textbox(elem_id="qf-message", show_label=False, placeholder="Type your message…", lines=1, max_lines=6),
        autofocus=False,
        fill_height=False,
        additional_inputs_accordion="Prompt and temperature",
        api_name="chat",
        title=None,
        description=None,
        additional_inputs=[system_prompt, temperature],
        cache_examples=False,
        save_history=False,
        examples=EXAMPLE_TEMPLATES,
        example_labels=["Batman", "Poetry", "Incident response"],
    )
    chat_events = [demo.fns[key].get_config() for key in demo.fns if key not in chat_event_start and demo.fns[key].queue]
    new_chat_btn.click(
        new_chat, inputs=None,
        outputs=[chat.chatbot, chat.chatbot_state, chat.chatbot_value,
                 chat.saved_input, chat.textbox, chat.conversation_id, chat.api_response],
        cancels=chat_events, queue=False, api_name="new_chat",
    )
    gr.Markdown('[Website](https://dumbbutt.tech) · [Hugging Face](https://huggingface.co/dumbbutt0) · [X](https://x.com/dumbbutt0)', elem_id="qf-footer")
    demo.unload(release_session)


# --------------------------------------------------
# STEP 9: LAUNCH
# --------------------------------------------------

if __name__ == "__main__":
    import uvicorn
    threading.Thread(target=expire_sessions, daemon=True).start()
    uvicorn.run(create_server(), host="0.0.0.0", port=int(os.environ.get("GRADIO_SERVER_PORT", "7860")), access_log=False)