from __future__ import annotations

import base64
import io
import json
import os
import secrets
import time
import uuid

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces
import gradio as gr
import httpx
import torch
from fastapi import Request
from fastapi.responses import JSONResponse, StreamingResponse
from gradio import Server
from PIL import Image
from peft import PeftConfig, PeftModel
from transformers import AutoProcessor, Gemma4ForConditionalGeneration

BASE_ID = os.environ.get(
    "BASE_MODEL_ID", "llmfan46/gemma-4-E4B-it-ultra-uncensored-heretic"
)
ADAPTER_ID = os.environ.get("ADAPTER_ID", "deathafteryou/gemma-4-hacking")
MODEL_ALIAS = os.environ.get("MODEL_ALIAS", "gemma-4-hacking")
API_KEY = os.environ.get("API_KEY", "")
MAX_NEW_TOKENS = int(os.environ.get("MAX_NEW_TOKENS", "2048"))

# The adapter ships target_modules as bare leaf names ("q_proj", ...). In this
# architecture the vision and audio projections are Gemma4ClippableLinear
# wrappers around an inner nn.Linear, so a bare leaf matches the wrapper and
# PEFT refuses to inject ("Target module Gemma4ClippableLinear ... is not
# supported"). Match the inner .linear of those wrappers explicitly, and keep
# bare leaves for the text tower, which uses plain nn.Linear.
ADAPTER_TARGET_MODULES = (
    r"(.*\.language_model\..*\.(q_proj|k_proj|v_proj|o_proj|gate_proj|up_proj|down_proj))"
    r"|(.*\.(q_proj|k_proj|v_proj|o_proj|gate_proj|up_proj|down_proj)\.linear)"
)

model = Gemma4ForConditionalGeneration.from_pretrained(
    BASE_ID,
    dtype=torch.bfloat16,
    low_cpu_mem_usage=True,
)
adapter_config = PeftConfig.from_pretrained(ADAPTER_ID)
adapter_config.target_modules = ADAPTER_TARGET_MODULES
# ZeroGPU reports torch.cuda.is_available() as True outside a @spaces.GPU call, so
# PEFT's infer_device() picks "cuda" and safetensors then tries a real CUDA
# allocation ("No CUDA GPUs are available"). Load the adapter on CPU; the
# following .to("cuda") is handled by ZeroGPU's deferred device aliases.
model = PeftModel.from_pretrained(
    model, ADAPTER_ID, config=adapter_config, torch_device="cpu"
)
model = model.to("cuda").eval()
processor = AutoProcessor.from_pretrained(BASE_ID)


def load_image(source):
    if isinstance(source, dict):
        source = source.get("url")
    if not isinstance(source, str) or not source:
        raise ValueError("image source is empty")
    if source.startswith("data:"):
        _, encoded = source.split(",", 1)
        raw = base64.b64decode(encoded)
    else:
        response = httpx.get(source, follow_redirects=True, timeout=30)
        response.raise_for_status()
        raw = response.content
    return Image.open(io.BytesIO(raw)).convert("RGB")


def normalize_messages(messages):
    normalized = []
    images = []
    for message in messages or []:
        if not isinstance(message, dict):
            continue
        role = message.get("role") or "user"
        content = message.get("content")
        if isinstance(content, str):
            normalized.append({"role": role, "content": content})
            continue
        if not isinstance(content, list):
            normalized.append({"role": role, "content": ""})
            continue
        parts = []
        for item in content:
            if not isinstance(item, dict):
                continue
            kind = item.get("type")
            if kind == "text":
                parts.append({"type": "text", "text": item.get("text") or ""})
            elif kind in ("image", "image_url"):
                source = item.get("image_url") or item.get("image")
                images.append(load_image(source))
                parts.append({"type": "image"})
        if not parts:
            parts = [{"type": "text", "text": ""}]
        normalized.append({"role": role, "content": parts})
    return normalized, images


def estimate_duration(messages=None, images=None, temperature=0.0, max_tokens=256, **kwargs):
    return max(30, min(120, 20 + int(max_tokens) // 2))


@spaces.GPU(duration=estimate_duration)
def generate(messages, images, temperature=0.0, max_tokens=256):
    text = processor.apply_chat_template(
        messages, add_generation_prompt=True, tokenize=False
    )
    arguments = {"text": text, "return_tensors": "pt"}
    if images:
        arguments["images"] = images
    inputs = processor(**arguments)
    inputs = {
        key: (value.to(model.device) if hasattr(value, "to") else value)
        for key, value in inputs.items()
    }
    with torch.inference_mode():
        output = model.generate(
            **inputs,
            max_new_tokens=int(max_tokens),
            do_sample=float(temperature) > 0,
            temperature=max(float(temperature), 1e-5),
            use_cache=True,
        )
    prompt_tokens = int(inputs["input_ids"].shape[1])
    generated = output[:, prompt_tokens:]
    content = processor.batch_decode(generated, skip_special_tokens=True)[0]
    return {
        "text": content,
        "prompt_tokens": prompt_tokens,
        "completion_tokens": int(generated.shape[1]),
    }


def extract_token(request: Request) -> str:
    # X-API-Key takes priority: a private Space consumes the Authorization header
    # for the Hub's own owner-only gate, so the API key must be able to travel in
    # its own header while Authorization carries the Hugging Face token.
    api_key_header = request.headers.get("x-api-key", "").strip()
    if api_key_header:
        return api_key_header
    header = request.headers.get("authorization", "")
    if header[:7].lower() == "bearer ":
        return header[7:].strip()
    return ""


def authorized(request: Request) -> bool:
    if not API_KEY:
        return False
    token = extract_token(request)
    if not token:
        return False
    return secrets.compare_digest(token, API_KEY)


def unauthorized() -> JSONResponse:
    return JSONResponse(
        {
            "error": {
                "message": "Invalid or missing API key",
                "type": "invalid_request_error",
                "code": "unauthorized",
            }
        },
        status_code=401,
        headers={"WWW-Authenticate": "Bearer"},
    )


app = Server(title="Gemma 4 Hacking Secure API")


@app.get("/healthz")
def healthz():
    return {"status": "ok", "ready": True, "model": MODEL_ALIAS}


@app.get("/v1/models")
def list_models(request: Request):
    if not authorized(request):
        return unauthorized()
    return {
        "object": "list",
        "data": [
            {
                "id": MODEL_ALIAS,
                "object": "model",
                "created": 0,
                "owned_by": "local",
            }
        ],
    }


def completion_payload(text, prompt_tokens, completion_tokens, model_name):
    return {
        "id": f"chatcmpl-{uuid.uuid4().hex}",
        "object": "chat.completion",
        "created": int(time.time()),
        "model": model_name,
        "choices": [
            {
                "index": 0,
                "message": {"role": "assistant", "content": text},
                "finish_reason": "stop",
            }
        ],
        "usage": {
            "prompt_tokens": prompt_tokens,
            "completion_tokens": completion_tokens,
            "total_tokens": prompt_tokens + completion_tokens,
        },
    }


def stream_chunks(payload):
    choice = payload["choices"][0]
    content = choice["message"]["content"]
    base = {
        "id": payload["id"],
        "object": "chat.completion.chunk",
        "created": payload["created"],
        "model": payload["model"],
    }
    for index in range(0, len(content), 24):
        chunk = dict(base)
        chunk["choices"] = [
            {
                "index": 0,
                "delta": {"content": content[index:index + 24]},
                "finish_reason": None,
            }
        ]
        yield f"data: {json.dumps(chunk)}\n\n"
    final = dict(base)
    final["choices"] = [{"index": 0, "delta": {}, "finish_reason": "stop"}]
    yield f"data: {json.dumps(final)}\n\n"
    yield "data: [DONE]\n\n"


@app.post("/v1/chat/completions")
def chat_completions(request: Request, payload: dict):
    if not authorized(request):
        return unauthorized()
    try:
        messages, images = normalize_messages(payload.get("messages"))
    except Exception as exc:
        return JSONResponse(
            {"error": {"message": str(exc), "type": "invalid_request_error"}},
            status_code=400,
        )
    max_tokens = payload.get("max_completion_tokens") or payload.get("max_tokens") or 256
    max_tokens = max(1, min(int(max_tokens), MAX_NEW_TOKENS))
    temperature = float(payload.get("temperature") or 0.0)
    try:
        result = generate(messages, images, temperature, max_tokens)
    except Exception as exc:
        return JSONResponse(
            {"error": {"message": str(exc), "type": "server_error"}},
            status_code=500,
        )
    body = completion_payload(
        result["text"],
        result["prompt_tokens"],
        result["completion_tokens"],
        payload.get("model") or MODEL_ALIAS,
    )
    if payload.get("stream"):
        return StreamingResponse(stream_chunks(body), media_type="text/event-stream")
    return JSONResponse(body)


@app.post("/v1/completions")
def completions(request: Request, payload: dict):
    if not authorized(request):
        return unauthorized()
    prompt = payload.get("prompt") or ""
    if isinstance(prompt, list):
        prompt = "\n".join(str(item) for item in prompt)
    max_tokens = payload.get("max_tokens") or 256
    max_tokens = max(1, min(int(max_tokens), MAX_NEW_TOKENS))
    temperature = float(payload.get("temperature") or 0.0)
    try:
        result = generate(
            [{"role": "user", "content": str(prompt)}], [], temperature, max_tokens
        )
    except Exception as exc:
        return JSONResponse(
            {"error": {"message": str(exc), "type": "server_error"}},
            status_code=500,
        )
    return JSONResponse(
        {
            "id": f"cmpl-{uuid.uuid4().hex}",
            "object": "text_completion",
            "created": int(time.time()),
            "model": payload.get("model") or MODEL_ALIAS,
            "choices": [
                {
                    "index": 0,
                    "text": result["text"],
                    "finish_reason": "stop",
                }
            ],
            "usage": {
                "prompt_tokens": result["prompt_tokens"],
                "completion_tokens": result["completion_tokens"],
                "total_tokens": result["prompt_tokens"] + result["completion_tokens"],
            },
        }
    )


with gr.Blocks(title="Gemma 4 Hacking Secure API") as demo:
    gr.Markdown(
        "# Gemma 4 Hacking Secure API\n\n"
        "`deathafteryou/gemma-4-hacking` served on ZeroGPU as an "
        "OpenAI-compatible API.\n\n"
        "All `/v1` routes require `Authorization: Bearer <API_KEY>` or "
        "`X-API-Key: <API_KEY>`. `/healthz` is the only public route.\n\n"
        "`POST /v1/chat/completions` accepts text and image content parts. "
        "Use the `openai` Python client with this Space URL plus `/v1`."
    )
    gr.Markdown(
        "The API key is stored in the Space secret `API_KEY`. Keep the Space "
        "URL and key private."
    )


if __name__ == "__main__":
    demo.launch(show_error=True, _app=app, ssr_mode=False)