#!/usr/bin/env python3
"""StarVector-1B HTTP server for HF Inference Endpoints (and local docker).

One contract, wherever it runs (PROPOSAL.md §5.1 L2 / §11):
    POST /generate  {"image_b64": "<png b64>", "max_tokens": 4096, "num_beams": 2}
                    -> {"svg": "<svg ...>", "elapsed_s": 12.3}
    GET  /health    -> {"ok": true, "model_loaded": bool, "device": "..."}

License notes: see LICENSE/NOTICE in this directory. Model weights are
self-contained in the starvector safetensors — we never touch the gated
bigcode/starcoderbase-1b repo (same trick as the estate bridge).
"""
import base64
import importlib.util
import io
import os
import sys
import textwrap
import time

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel

MODEL_NAME = os.environ.get("STARVECTOR_MODEL", "starvector/starvector-1b-im2svg")
MAX_TOKENS_DEFAULT = int(os.environ.get("STARVECTOR_MAX_TOKENS", "4096"))
PORT = int(os.environ.get("PORT", "8080"))

app = FastAPI(title="IconEdge StarVector bridge", version="0.1.0")
_state = {"model": None, "device": "cpu", "loaded": False}


def _patch_starvector() -> None:
    """Same patch as IconTreasury's bridge: load tokenizer+LLM weights from the
    starvector repo itself instead of the gated bigcode/starcoderbase-1b repo."""
    spec = importlib.util.find_spec("starvector")
    if spec is None or not spec.submodule_search_locations:
        raise RuntimeError("starvector package not installed")
    pkg_dir = list(spec.submodule_search_locations)[0]
    patch_path = os.path.join(pkg_dir, "model", "llm", "starcoder_patched.py")
    if os.path.exists(patch_path):
        return
    os.makedirs(os.path.dirname(patch_path), exist_ok=True)
    with open(patch_path, "w") as f:
        f.write(textwrap.dedent('''\
            """Patched StarCoderModel — loads tokenizer+weights from starvector repo
            instead of the gated bigcode/starcoderbase-1b repo."""
            import torch.nn as nn
            from transformers import AutoTokenizer
            from starvector.model.gpt_bigcode.configuration_gpt_bigcode import GPTBigCodeConfig
            from starvector.model.gpt_bigcode.modeling_gpt_bigcode import GPTBigCodeForCausalLM

            class StarCoderModelPatched(nn.Module):
                def __init__(self, config, **kwargs):
                    super().__init__()
                    repo = config._name_or_path if hasattr(config, "_name_or_path") else "starvector/starvector-1b-im2svg"
                    self.init_tokenizer(repo)
                    self.max_length = config.max_length
                    dropout = config.dropout if hasattr(config, "dropout") else 0.1
                    model_config = GPTBigCodeConfig(
                        vocab_size=config.vocab_size,
                        n_positions=config.max_position_embeddings,
                        n_embd=config.hidden_size,
                        n_layer=config.num_hidden_layers,
                        n_head=config.num_attention_heads,
                        n_inner=config.hidden_size * 4,
                        activation_function="gelu_pytorch_tanh",
                        resid_pdrop=dropout, embd_pdrop=dropout, attn_pdrop=dropout,
                        layer_norm_epsilon=1e-5, scale_attn_weights=True,
                        use_cache=config.use_cache,
                        bos_token_id=self.tokenizer.bos_token_id,
                        eos_token_id=self.tokenizer.eos_token_id,
                        pad_token_id=self.tokenizer.pad_token_id if self.tokenizer.pad_token_id else self.tokenizer.eos_token_id,
                        multi_query=config.multi_query if hasattr(config, "multi_query") else True,
                    )
                    model = GPTBigCodeForCausalLM(model_config)
                    model.resize_token_embeddings(len(self.tokenizer))
                    self.transformer = model
                    self.prompt = "<svg"

                def init_tokenizer(self, model_name):
                    self.tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False)
                    if self.tokenizer.eos_token_id is None:
                        self.tokenizer.add_special_tokens({"eos_token": "[EOS]"})
                    if self.tokenizer.pad_token_id is None:
                        self.tokenizer.add_special_tokens({"pad_token": "[PAD]"})
                    self.svg_start_token = "<svg-start>"
                    self.image_start_token = "<image-start>"
                    self.text_start_token = "<caption-start>"
                    self.tokenizer.add_tokens([self.svg_start_token, self.image_start_token, self.text_start_token])
                    self.svg_start_token_id = self.tokenizer.encode(self.svg_start_token)[0]
            '''))


def _load_model() -> None:
    import torch

    _patch_starvector()
    import starvector.model.models.starvector_v1 as v1
    from starvector.model.llm.starcoder_patched import StarCoderModelPatched
    v1.StarVectorStarCoder._get_svg_transformer = (
        lambda self, config, **kwargs: StarCoderModelPatched(config, **kwargs)
    )
    from starvector.model.starvector_arch import StarVectorForCausalLM

    device = "cuda" if torch.cuda.is_available() else "cpu"
    dtype = torch.float16
    print(f"Loading {MODEL_NAME} on {device}...", flush=True)
    t0 = time.time()
    model = StarVectorForCausalLM.from_pretrained(
        MODEL_NAME, torch_dtype=dtype, trust_remote_code=True,
    )
    model.config.use_flash_attn = False
    model.eval()
    model = model.to(device)
    model.model_precision = dtype
    _state.update(model=model, device=device, loaded=True)
    print(f"Model loaded in {time.time() - t0:.1f}s", flush=True)


class GenerateRequest(BaseModel):
    image_b64: str
    max_tokens: int = MAX_TOKENS_DEFAULT
    num_beams: int = 2


@app.get("/health")
def health():
    return {"ok": True, "model_loaded": _state["loaded"], "device": _state["device"]}


@app.post("/generate")
def generate(req: GenerateRequest):
    if not _state["loaded"]:
        raise HTTPException(503, "model still loading — retry shortly")
    import torch
    from PIL import Image

    try:
        png_bytes = base64.b64decode(req.image_b64)
    except Exception as e:
        raise HTTPException(400, f"image_b64 is not valid base64: {e}")

    image = Image.open(io.BytesIO(png_bytes)).convert("RGB")
    model = _state["model"]
    image_tensor = model.process_images([image])[0].to(_state["device"], dtype=torch.float16)

    t0 = time.time()
    with torch.no_grad():
        raw_svgs = model.generate_im2svg(
            {"image": image_tensor}, max_length=req.max_tokens,
            num_beams=req.num_beams, temperature=1.0,
        )
    return {"svg": raw_svgs[0], "elapsed_s": round(time.time() - t0, 2)}


@app.on_event("startup")
def startup():
    if os.environ.get("LOAD_ON_STARTUP", "1") == "1":
        _load_model()


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=PORT)