from transformers import pipeline
import gradio as gr
import spaces


pipe = pipeline("text-generation", model="markrodrigo/Llama-3.2-3B-Instruct-Spatial-SQL-1.1", device_map="auto")

# The Alpaca instruction prompt format
ALPACA_TEMPLATE = """<|begin_of_text|><|start_header_id|>system<|end_header_id|> You are a helpful assistant. You are an expert at PostGIS and Postgresql and SQL and psql. <|eot_id|><|start_header_id|>user<|end_header_id|>

### Instruction: Write a PostGIS SQL statement for the following.
{instruction}

### Input:
{input}

### Response:
<|eot_id|><|start_header_id|>assistant<|end_header_id|>
"""

# Define your list of pre-set example prompts
PRESET_EXAMPLES = [
    "What is the area for the polygon? : 'Polygon ((-3.7515154 40.3855551, -3.7514972 40.3856581, -3.7507005 40.3855767, -3.7507167 40.3854722, -3.7515154 40.3855551))'",
    "What is the centroid for the polygon? : 'Polygon ((-3.6934636 40.4808785, -3.6933352 40.4811486, -3.6930125 40.4810598, -3.693141 40.4807897, -3.6934636 40.4808785))'",
    "What is the thousand meter buffer for the following point? : 'Point(-8.7522658 41.3862664)'",
    "How long is the line? : 'LINESTRING (-3.6976693 40.4263178, -3.6986082 40.4258729)'",
    "How far apart is the point and line? : 'Point(-109.87549823 38.60574249)' 'LineString(-109.24324628 38.76349931, -109.4821773 38.6875815)'"
]


@spaces.GPU
def respond(user_message, chat_history):
    chat_history = chat_history or []

    if not user_message or not user_message.strip():
        return chat_history, ""

    # Modern format for Gradio 5/6
    chat_history.append({"role": "user", "content": user_message})
    chat_history.append({"role": "assistant", "content": None})

    prompt = ALPACA_TEMPLATE.format(instruction=user_message, input="")

    sequences = pipe(
        prompt,
        max_new_tokens=256,
        return_full_text=False,
        temperature=0.4,
        top_k=100,
        do_sample=True,
    )

    bot_response = sequences[0]["generated_text"].strip()
    chat_history[-1]["content"] = bot_response

    return chat_history, ""



with gr.Blocks(title="Text to PostGIS Postgresql via Llama 3.2") as demo:
    gr.Markdown("# Natural Language to Spatial SQL.\n###  Convert natural language and spatial primitives to PostGIS with Llama 3.2")
    chatbot = gr.Chatbot(
        label="Chat",
        height=400,
        # type="messages",
        # show_copy_button=True,
    )

    with gr.Row():
        with gr.Column(scale=5):
            msg = gr.Textbox(
                placeholder="Natural Language : WKT format",
                lines=2,
                container=False
            )
        with gr.Column(scale=1, min_width=100):
            submit_btn = gr.Button("Submit", variant="primary")

    gr.Markdown("### Quick Examples")
    gr.Examples(
        examples=PRESET_EXAMPLES,
        inputs=msg,
        label="Click an example → then click Submit"
    )

    submit_btn.click(fn=respond, inputs=[msg, chatbot], outputs=[chatbot, msg])
    msg.submit(fn=respond, inputs=[msg, chatbot], outputs=[chatbot, msg])

    clear_btn = gr.Button("Clear Chat")
    clear_btn.click(lambda: ([], ""), outputs=[chatbot, msg])


if __name__ == "__main__":
    print("Gradio version:", gr.__version__)
    demo.launch()

'''from transformers import pipeline
import gradio as gr
import spaces

# Initialize the pipeline with an Alpaca-tuned model
pipe = pipeline("text-generation", model="markrodrigo/Llama-3.2-3B-Instruct-Spatial-SQL-1.1", device_map="auto")

# The Alpaca instruction prompt format
ALPACA_TEMPLATE = """<|begin_of_text|><|start_header_id|>system<|end_header_id|> You are a helpful assistant. You are an expert at PostGIS and Postgresql and SQL and psql. <|eot_id|><|start_header_id|>user<|end_header_id|>

### Instruction: Write a PostGIS SQL statement for the following.
{instruction}

### Input:
{input}

### Response:
<|eot_id|><|start_header_id|>assistant<|end_header_id|>
"""

# Define your list of pre-set example prompts
PRESET_EXAMPLES = [
    "What is the area for the polygon? : 'Polygon ((-3.7515154 40.3855551, -3.7514972 40.3856581, -3.7507005 40.3855767, -3.7507167 40.3854722, -3.7515154 40.3855551))'",
    "What is the centroid for the polygon? : 'Polygon ((-3.6934636 40.4808785, -3.6933352 40.4811486, -3.6930125 40.4810598, -3.693141 40.4807897, -3.6934636 40.4808785))'",
    "What is the thousand meter buffer for the following point? : 'Point(-8.7522658 41.3862664)'",
    "How long is the line? : 'LINESTRING (-3.6976693 40.4263178, -3.6986082 40.4258729)'",
    "How far apart is the point and line? : 'Point(-109.87549823 38.60574249)' 'LineString(-109.24324628 38.76349931, -109.4821773 38.6875815)'"
]


@spaces.GPU
def chat_predict(message, history):
    print("gradio " + gr.__version__)
    # For a chat interface, we treat the latest user message as the instruction
    # and leave the input field empty for this example.
    prompt = ALPACA_TEMPLATE.format(instruction=message, input="")
    
    # Generate text (adjust max_new_tokens as needed)
    sequences = pipe(
        prompt,
        max_new_tokens=128,
        return_full_text=False,
        temperature=0.4,
        top_k=100,
        # top_p=0.9,
    )
    
    # Extract and return the generated text
    response = sequences[0]['generated_text'].strip()
    return response

# Create the Gradio ChatInterface
demo = gr.ChatInterface(
    fn=chat_predict,
    title="Text to PostGIS SQL via Llama 3.2 - Primary Functions",
    description="LLama 3.2 Spatial - Text to PostGIS SQL",
)

if __name__ == "__main__":
    print("gradio " + gr.__version__)
    demo.launch()
'''