import gradio as gr
import spaces
import json
import hashlib
import logging
from functools import lru_cache
from transformers import pipeline

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

_pipeline_cache = {"hash": None, "pipe": None, "config": None}

def _hash_config(cfg: dict) -> str:
    return hashlib.sha256(json.dumps(cfg, sort_keys=True).encode()).hexdigest()[:16]

@lru_cache(maxsize=3)
def _load_pipeline(hash_key: str, cfg_json: str):
    cfg = json.loads(cfg_json)
    logger.info(f"Loading pipeline: {cfg}")
    pipe = pipeline(**cfg)
    logger.info("Loaded.")
    return pipe

def get_pipe(cfg: dict):
    h = _hash_config(cfg)
    if _pipeline_cache["hash"] == h and _pipeline_cache["pipe"] is not None:
        return _pipeline_cache["pipe"], False
    pipe = _load_pipeline(h, json.dumps(cfg, sort_keys=True))
    _pipeline_cache.update({"hash": h, "pipe": pipe, "config": cfg})
    return pipe, True

@spaces.GPU(duration=40)
def inference(pipeline_config, inputs, inference_kwargs=None):
    if inference_kwargs is None:
        inference_kwargs = {}
    try:
        pconf = json.loads(pipeline_config) if isinstance(pipeline_config, str) else pipeline_config
    except Exception as e:
        return json.dumps({"error": f"Bad pipeline_config: {e}"})
    try:
        ikw = json.loads(inference_kwargs) if isinstance(inference_kwargs, str) else inference_kwargs
    except Exception as e:
        return json.dumps({"error": f"Bad inference_kwargs: {e}"})

    try:
        pipe, reloaded = get_pipe(pconf)
    except Exception as e:
        return json.dumps({"error": f"Pipeline load failed: {e}"})

    try:
        raw = json.loads(inputs) if isinstance(inputs, str) and inputs.strip().startswith(("[", "{")) else inputs
    except:
        raw = inputs

    try:
        result = pipe(raw, **ikw)
        return json.dumps({
            "success": True,
            "reloaded": reloaded,
            "result": result,
        }, default=str, indent=2)
    except Exception as e:
        return json.dumps({"error": f"Inference failed: {e}"})

# Example configs
EX1 = {
    "pipeline_config": {"task": "text-generation", "model": "HuggingFaceTB/SmolLM2-135M-Instruct"},
    "inputs": "The future of AI is",
    "inference_kwargs": {"max_new_tokens": 50},
}
EX2 = {
    "pipeline_config": {"task": "zero-shot-classification", "model": "facebook/bart-large-mnli"},
    "inputs": "This is a contract about data privacy and user rights.",
    "inference_kwargs": {"candidate_labels": ["legal", "finance", "technology", "sports"]},
}
EX3 = {
    "pipeline_config": {"task": "token-classification", "model": "openai/privacy-filter"},
    "inputs": "My name is Alice Smith",
    "inference_kwargs": {},
}

def run_example(ex):
    result = inference(ex["pipeline_config"], ex["inputs"], ex["inference_kwargs"])
    return ex["pipeline_config"], ex["inputs"], ex["inference_kwargs"], result

with gr.Blocks(title="Dynamic Transformers Pipeline API") as demo:
    gr.Markdown(
        "# 🚀 Dynamic Transformers Pipeline API\n"
        "[Continue AI conversation to edit](https://huggingface.co/chat/conversation/6a751b859a0b83e84c4abd74)\n\n"
        "Zero-GPU. Pass any `transformers.pipeline` config via JSON."
    )

    with gr.Row():
        with gr.Column(scale=1):
            pcfg = gr.JSON(label="pipeline_config", value=EX1["pipeline_config"])
            inp = gr.Textbox(label="inputs", value=EX1["inputs"], lines=3)
            ikw = gr.JSON(label="inference_kwargs", value=EX1["inference_kwargs"])
            btn = gr.Button("▶️ Run Inference", variant="primary")

        with gr.Column(scale=1):
            out = gr.JSON(label="output")

    btn.click(inference, [pcfg, inp, ikw], out)

    gr.Markdown("## 📝 Examples (click to populate & run)")
    with gr.Row():
        with gr.Column():
            gr.Markdown("**Text Generation**")
            ex1_btn = gr.Button("Run: SmolLM2 text-gen", size="sm")
            ex1_out = gr.JSON(label="result")

        with gr.Column():
            gr.Markdown("**Zero-Shot Classification**")
            ex2_btn = gr.Button("Run: zero-shot", size="sm")
            ex2_out = gr.JSON(label="result")

        with gr.Column():
            gr.Markdown("**Token Classification (Privacy Filter)**")
            ex3_btn = gr.Button("Run: privacy-filter", size="sm")
            ex3_out = gr.JSON(label="result")

    ex1_btn.click(
        fn=lambda: run_example(EX1),
        inputs=None,
        outputs=[pcfg, inp, ikw, ex1_out],
    )
    ex2_btn.click(
        fn=lambda: run_example(EX2),
        inputs=None,
        outputs=[pcfg, inp, ikw, ex2_out],
    )
    ex3_btn.click(
        fn=lambda: run_example(EX3),
        inputs=None,
        outputs=[pcfg, inp, ikw, ex3_out],
    )

    gr.Markdown("## API Example")
    gr.Code("""from gradio_client import Client

client = Client("DoctorSlimm/dynamic-transformers-api")
print(client.predict(
    pipeline_config={"task": "text-generation", "model": "HuggingFaceTB/SmolLM2-135M-Instruct"},
    inputs="The future of AI is",
    inference_kwargs={"max_new_tokens": 50},
    api_name="/inference"
))""", language="python")

demo.launch()