import gradio as gr
import json
import torch
import os
from transformers import AutoModelForCausalLM, AutoTokenizer
import spaces

title = """        # 🙋🏻‍♂️Welcome to 🌟Tonic's 🌊 Osmosis Structure - Text to JSON Converter
"""
description = """       
        Convert unstructured text into well-formatted JSON using the Osmosis Structure 0.6B model.
        This model is specifically trained for structured data extraction and format conversion.
        
        ### ℹ️ About Osmosis Structure
        
        - **Model**: Osmosis Structure 0.6B parameters
        - **Architecture**: Qwen3 (specialized for structured data)
        - **Purpose**: Converting unstructured text to structured JSON format
        - **Optimizations**: Fine-tuned for data extraction and format conversion tasks
        - **Access**: Requires HF authentication token for gated repository
        
        The model automatically identifies key information in your text and organizes it into logical JSON structures.
        """
joinus = """
## Join us :
🌟TeamTonic🌟 is always making cool demos! Join our active builder's 🛠️community 👻 [![Join us on Discord](https://img.shields.io/discord/1109943800132010065?label=Discord&logo=discord&style=flat-square)](https://discord.gg/qdfnvSPcqP) On 🤗Huggingface:[MultiTransformer](https://huggingface.co/MultiTransformer) On 🌐Github: [Tonic-AI](https://github.com/tonic-ai) & contribute to🌟 [MultiTonic](https://github.com/MultiTonic)🤗Big thanks to Yuvi Sharma and all the folks at huggingface for the community grant 🤗
"""

# Model configuration
MODEL_NAME = "osmosis-ai/Osmosis-Structure-0.6B"

# Global variables to store the model and tokenizer
model = None
tokenizer = None

def load_model():
    """Load the Osmosis Structure model and tokenizer with HF token for gated repos"""
    global model, tokenizer
    
    try:
        print("Loading Osmosis Structure model...")
        
        # Get HF token from environment variables
        hf_token = os.environ.get("HF_KEY")
        if not hf_token:
            print("⚠️ Warning: HF_KEY not found in environment variables")
            print("Attempting to load without token...")
            hf_token = None
        else:
            print("✅ HF token found, accessing gated repository...")
        
        # Load tokenizer with token
        print("Loading tokenizer...")
        tokenizer = AutoTokenizer.from_pretrained(
            MODEL_NAME,
            trust_remote_code=True,
            token=hf_token
            )
        
        print("Loading model...")
        # Load model with token
        model = AutoModelForCausalLM.from_pretrained(
            MODEL_NAME,
            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
            device_map="auto" if torch.cuda.is_available() else None,
            trust_remote_code=True,
            token=hf_token,
            use_auth_token=hf_token  # Backward compatibility
        )
        
        print("✅ Osmosis Structure model loaded successfully!")
        return True
        
    except Exception as e:
        error_msg = f"❌ Error loading model: {e}"
        print(error_msg)
        
        # Provide helpful error messages for common issues
        if "401" in str(e) or "authentication" in str(e).lower():
            print("💡 This appears to be an authentication error.")
            print("Please ensure:")
            print("1. HF_KEY is set correctly in your Space secrets")
            print("2. Your token has access to the gated repository")
            print("3. You have accepted the model's license agreement")
        elif "404" in str(e) or "not found" in str(e).lower():
            print("💡 Model repository not found.")
            print("Please check if the model name is correct and accessible")
        
        return False

@spaces.GPU
def text_to_json(input_text, max_tokens=512, temperature=0.6, top_p=0.95, top_k=20):
    """Convert plain text to structured JSON using Osmosis Structure model"""
    global model, tokenizer
    
    if model is None or tokenizer is None:
        return "❌ Model not loaded. Please check the console for loading errors."
    
    try:
        # Create a structured prompt for JSON conversion
        messages = [
            {
                "role": "system",
                "content": "You are a helpful assistant that converts unstructured text into well-formatted JSON. Extract key information and organize it into a logical, structured format. Always respond with valid JSON."
            },
            {
                "role": "user", 
                "content": f"Convert this text to JSON format:\n\n{input_text}"
            }
        ]
        
        # Apply chat template
        formatted_prompt = tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True
        )
        
        # Tokenize the input
        inputs = tokenizer(
            formatted_prompt,
            return_tensors="pt",
            truncation=True,
            max_length=2048
        )
        
        # Move to device if using GPU
        if torch.cuda.is_available():
            inputs = {k: v.to(model.device) for k, v in inputs.items()}
        
        # Generation parameters based on model config
        generation_config = {
            "max_new_tokens": max_tokens,
            "temperature": temperature,
            "top_p": top_p,
            "top_k": top_k,
            "do_sample": True,
            "pad_token_id": tokenizer.pad_token_id,
            "eos_token_id": tokenizer.eos_token_id,
            "repetition_penalty": 1.1,
        }
        
        # Generate response
        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                **generation_config
            )
        
        # Decode the response
        generated_tokens = outputs[0][len(inputs["input_ids"][0]):]
        generated_text = tokenizer.decode(generated_tokens, skip_special_tokens=True)
        
        # Clean up the response
        generated_text = generated_text.strip()
        
        # Try to extract JSON from the response
        json_start = generated_text.find('{')
        json_end = generated_text.rfind('}')
        
        if json_start != -1 and json_end != -1 and json_end > json_start:
            json_text = generated_text[json_start:json_end+1]
        else:
            # If no clear JSON boundaries, try to clean the whole response
            json_text = generated_text
            
            # Remove common prefixes
            prefixes_to_remove = ["```json", "```", "Here's the JSON:", "JSON:", "```json\n"]
            for prefix in prefixes_to_remove:
                if json_text.startswith(prefix):
                    json_text = json_text[len(prefix):].strip()
            
            # Remove common suffixes
            suffixes_to_remove = ["```", "\n```"]
            for suffix in suffixes_to_remove:
                if json_text.endswith(suffix):
                    json_text = json_text[:-len(suffix)].strip()
        
        # Validate and format JSON
        try:
            parsed_json = json.loads(json_text)
            return json.dumps(parsed_json, indent=2, ensure_ascii=False)
        except json.JSONDecodeError:
            # If still not valid JSON, return the cleaned text with a note
            return f"Generated response (may need manual cleanup):\n\n{json_text}"
            
    except Exception as e:
        return f"❌ Error generating JSON: {str(e)}"

def create_demo():
    # Fixed: Remove duplicate with gr.Blocks declaration
    with gr.Blocks(
        title=title,
        theme=gr.themes.Monochrome(),
        css="""
        .gradio-container {
            max-width: 1200px !important;
        }
        """
    ) as demo:
        # Header section
        # gr.Markdown(title)
        
        # Info section
        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown(description)
            with gr.Column(scale=1):
                gr.Markdown(joinus)
        
        with gr.Row():
            with gr.Column(scale=1):
                input_text = gr.Textbox(
                    label="📝 Input Text",
                    placeholder="Enter your unstructured text here...\n\nExample: 'John Smith is a 30-year-old software engineer from New York. He works at Tech Corp and has 5 years of experience in Python development.'",
                    lines=8,
                    max_lines=15
                )
                
                with gr.Accordion("⚙️ Generation Settings", open=False):
                    max_tokens = gr.Slider(
                        minimum=50,
                        maximum=1000,
                        value=512,
                        step=10,
                        label="Max Tokens",
                        info="Maximum number of tokens to generate"
                    )
                    
                    temperature = gr.Slider(
                        minimum=0.1,
                        maximum=1.0,
                        value=0.6,
                        step=0.1,
                        label="Temperature",
                        info="Controls randomness (lower = more focused)"
                    )
                    
                    top_p = gr.Slider(
                        minimum=0.1,
                        maximum=1.0,
                        value=0.95,
                        step=0.05,
                        label="Top-p",
                        info="Nucleus sampling parameter"
                    )
                    
                    top_k = gr.Slider(
                        minimum=1,
                        maximum=100,
                        value=20,
                        step=1,
                        label="Top-k",
                        info="Limits vocabulary for generation"
                    )
                
                convert_btn = gr.Button(
                    "🔄 Convert to JSON",
                    variant="primary",
                    size="lg"
                )
                
            with gr.Column(scale=1):
                output_json = gr.Textbox(
                    label="📋 Generated JSON",
                    lines=15,
                    max_lines=20,
                    interactive=False,
                    show_copy_button=True
                )
        
        # Example inputs
        gr.Markdown("### 📚 Example Inputs")
        examples = gr.Examples(
            examples=[
                ["John Smith is a 30-year-old software engineer from New York. He works at Tech Corp and has 5 years of experience in Python development. His email is john.smith@email.com and he graduated from MIT in 2018."],
                ["Order #12345 was placed on March 15, 2024. Customer: Sarah Johnson, Address: 123 Main St, Boston MA 02101. Items: 2x Laptop ($999 each), 1x Mouse ($25). Total: $2023. Status: Shipped via FedEx, tracking: 1234567890."],
                ["The conference will be held on June 10-12, 2024 at the Grand Hotel in San Francisco. Registration fee is $500 for early bird (before May 1) and $650 for regular registration. Contact info@conference.com for questions."],
                ["Product: Wireless Headphones Model XYZ-100. Price: $199.99. Features: Bluetooth 5.0, 30-hour battery, noise cancellation, wireless charging case. Colors available: Black, White, Blue. Warranty: 2 years. Rating: 4.5/5 stars (324 reviews)."]
            ],
            inputs=input_text,
            label="Click on any example to try it"
        )
        

        # Event handlers
        convert_btn.click(
            fn=text_to_json,
            inputs=[input_text, max_tokens, temperature, top_p, top_k],
            outputs=output_json,
            show_progress=True
        )
        
        # Allow Enter key to trigger conversion
        input_text.submit(
            fn=text_to_json,
            inputs=[input_text, max_tokens, temperature, top_p, top_k],
            outputs=output_json,
            show_progress=True
        )
    
    return demo

# Initialize the demo
if __name__ == "__main__":
    print("🌊 Initializing Osmosis Structure Demo...")
    
    # Check HF token availability
    hf_token = os.environ.get("HF_KEY")
    if hf_token:
        print("✅ HF_KEY found in environment")
    else:
        print("⚠️ HF_KEY not found - this may cause issues with gated repositories")
    
    # Load model at startup
    if load_model():
        print("🚀 Creating Gradio interface...")
        demo = create_demo()
        demo.launch(
            ssr_mode=False,
            mcp_server=True
        )
    else:
        print("❌ Failed to load model. Please check your HF_KEY and model access permissions.")