import torch
from transformers import pipeline
import gradio as gr
import json

# Initialize model pipeline
model_id = "dylanhogg/gnaf-structured-address-v0.1-75a1791-20250921-063650"

pipe = pipeline(
    "text-generation",
    model=model_id,
    dtype=torch.bfloat16,
    device_map="auto",
)


def parse_address(user_address: str) -> tuple[str, str]:
    """Parse address and return both raw response and JSON"""
    user_content = f"Translate a text address into structured json.\n{user_address}"
    messages = [
        {"role": "system", "content": "Translate a text address into structured json."},
        {"role": "user", "content": user_content},
    ]

    outputs = pipe(
        messages,
        max_new_tokens=256,
        do_sample=False,
    )

    response = outputs[0]["generated_text"]
    last_content = response[-1]["content"]

    # Try to extract JSON from the response
    try:
        # Find first '{' and last '}' to extract JSON
        start = last_content.find("{")
        end = last_content.rfind("}") + 1
        if start != -1 and end > start:
            json_str = last_content[start:end]
            parsed = json.loads(json_str)
            formatted_json = json.dumps(parsed, indent=2)
        else:
            formatted_json = "No JSON found in response"
    except Exception as e:
        formatted_json = f"Error parsing JSON: {str(e)}\n\nRaw output:\n{last_content}"

    return formatted_json, last_content


# Create Gradio interface
with gr.Blocks(title="Address Parser") as demo:
    gr.Markdown("# 🏠 Structured Address Parser")
    gr.Markdown("This model converts text addresses into structured JSON format.")

    with gr.Row():
        with gr.Column():
            input_text = gr.Textbox(
                label="Input Address", placeholder="Enter an address...", value="48a Pirrama Rd Pyrmont NSW 2009"
            )
            submit_btn = gr.Button("Parse Address", variant="primary")

        with gr.Column():
            json_output = gr.Textbox(label="Structured JSON", interactive=False, lines=10)
            raw_output = gr.Textbox(label="Raw Model Output", interactive=False, lines=5)

    # Examples
    gr.Examples(
        examples=[
            "48a Pirrama Rd Pyrmont NSW 2009",
            "Floor 3, 152-156 Clarence St, Sydney NSW 2000",
            "Aptt 16, 400 Bondi Rd, Bondi NSW 2026",
            "Unit 18/14-18 Flood St, Bondi, NSW 2026",
            "Lvl 15/333 George St Sydney NSW 2000",
            "Check out: 44 Ulm St, Maroubra NSW 2035",
            "44 Ulm St, Maroubra NSW 2035 is where it's at!",
        ],
        inputs=input_text,
    )

    # Handle events
    submit_btn.click(fn=parse_address, inputs=input_text, outputs=[json_output, raw_output])
    input_text.submit(fn=parse_address, inputs=input_text, outputs=[json_output, raw_output])

# Launch the app
if __name__ == "__main__":
    demo.launch()