import gradio as gr
import json
import requests
import os

# Use the Hugging Face Inference API instead of loading the model
API_URL = "/static-proxy?url=https%3A%2F%2Fapi-inference.huggingface.co%2Fmodels%2Fnumind%2FNuExtract-1.5"
headers = {"Authorization": f"Bearer {os.environ.get('HF_TOKEN', '')}"}

def test_function(template, text):
    print(f"Test function called with template: {template[:30]} and text: {text[:30]}")
    return "Button clicked successfully", "Function was called"

def extract_info(template, text):
    try:
        # Format prompt according to NuExtract-1.5 requirements
        prompt = f"<|input|>\n### Template:\n{template}\n### Text:\n{text}\n\n<|output|>"
        print(f"Processing with prompt: {prompt[:100]}...")
        
        # Call API instead of using local model
        payload = {
            "inputs": prompt,
            "parameters": {
                "max_new_tokens": 1000,
                "do_sample": False
            }
        }
        
        print("Calling API...")
        response = requests.post(API_URL, headers=headers, json=payload)
        
        if response.status_code != 200:
            print(f"API error: {response.status_code}, {response.text}")
            return f"❌ API Error: {response.status_code}", response.text
        
        # Process result
        result = response.json()
        
        # Handle different response formats
        if isinstance(result, list) and len(result) > 0:
            result_text = result[0].get("generated_text", "")
        else:
            result_text = str(result)
        
        # Split at output marker if present
        if "<|output|>" in result_text:
            json_text = result_text.split("<|output|>")[1].strip()
        else:
            json_text = result_text
        
        # Try to parse as JSON
        print("Parsing JSON...")
        try:
            extracted = json.loads(json_text)
            formatted = json.dumps(extracted, indent=2)
        except json.JSONDecodeError:
            print(f"JSON parsing failed. Raw output: {json_text[:100]}...")
            return "❌ JSON parsing error", json_text
            
        return "✅ Success", formatted
    except Exception as e:
        print(f"Error in extraction: {str(e)}")
        return f"❌ Error: {str(e)}", "{}"

# Create a simple interface
with gr.Blocks() as demo:
    gr.Markdown("# NuExtract-1.5 Extraction Tool")
    
    with gr.Row():
        with gr.Column():
            template = gr.Textbox(
                label="JSON Template", 
                value='{"name": "", "email": ""}',
                lines=5
            )
            text = gr.Textbox(
                label="Text to Extract From",
                value="Contact: John Smith (john@example.com)",
                lines=8
            )
            
            # Two buttons for testing
            test_btn = gr.Button("Test Click")
            extract_btn = gr.Button("Extract Information", variant="primary")
        
        with gr.Column():
            status = gr.Textbox(label="Status")
            output = gr.Textbox(label="Output", lines=10)
    
    # Connect both buttons
    test_btn.click(
        fn=test_function,
        inputs=[template, text],
        outputs=[status, output]
    )
    
    extract_btn.click(
        fn=extract_info,
        inputs=[template, text],
        outputs=[status, output]
    )

if __name__ == "__main__":
    demo.launch()