import gradio as gr
import json
from llama_cpp import Llama
import os
from huggingface_hub import hf_hub_download
from config import get_model_config, get_generation_config, get_recommended_model

# Global variable to store the model
llm = None

def load_model():
    """Load the llama.cpp model"""
    global llm
    try:
        print("Loading Osmosis Structure model...")
        
        # Get model info and config
        model_info = get_recommended_model()
        model_config = get_model_config()
        
        # Create models directory
        os.makedirs("./models", exist_ok=True)
        
        # Download the Osmosis model
        print(f"Downloading {model_info['name']} ({model_info['size']})...")
        model_path = hf_hub_download(
            repo_id=model_info['repo_id'],
            filename=model_info['filename'],
            cache_dir="./models",
            resume_download=True
        )
        
        print(f"Model downloaded to: {model_path}")
        print("Initializing llama.cpp...")
        
        # Initialize llama.cpp with the downloaded model
        llm = Llama(
            model_path=model_path,
            **model_config
        )
        
        print("✅ Osmosis Structure model loaded successfully!")
        return f"✅ Model loaded: {model_info['name']}\nPath: {model_path}\nDescription: {model_info['description']}"
        
    except Exception as e:
        error_msg = f"❌ Error loading model: {e}"
        print(error_msg)
        return error_msg

def text_to_json(input_text, max_tokens=512, temperature=0.7):
    """Convert plain text to structured JSON using llama.cpp"""
    global llm
    
    if llm is None:
        return "❌ Model not loaded. Please load the model first."
    
    try:
        # Create a structured prompt optimized for the Osmosis model
        prompt = f"""<|system|>
You are a helpful assistant that converts unstructured text into well-formatted JSON. Extract key information and organize it into a logical structure.

<|user|>
Convert this text to JSON format:

{input_text}

<|assistant|>
```json"""

        # Get generation config and override with user settings
        gen_config = get_generation_config()
        gen_config.update({
            "max_tokens": max_tokens,
            "temperature": temperature
        })
        
        # Generate response using llama.cpp
        response = llm(
            prompt,
            **gen_config,
            echo=False
        )
        
        generated_text = response['choices'][0]['text'].strip()
        
        # Clean up the response - remove markdown formatting if present
        if generated_text.startswith('```json'):
            generated_text = generated_text[7:]
        if generated_text.endswith('```'):
            generated_text = generated_text[:-3]
        generated_text = generated_text.strip()
        
        # Try to parse as JSON to validate
        try:
            parsed_json = json.loads(generated_text)
            return json.dumps(parsed_json, indent=2)
        except json.JSONDecodeError:
            # If not valid JSON, try to clean it up or return as is
            return f"Generated (may need cleanup):\n{generated_text}"
            
    except Exception as e:
        return f"❌ Error generating JSON: {str(e)}"

def demo_without_model(input_text):
    """Demo function that works without loading a model"""
    try:
        # Simple rule-based JSON conversion for demonstration
        words = input_text.strip().split()
        
        # Create a basic JSON structure
        result = {
            "input_text": input_text,
            "word_count": len(words),
            "words": words,
            "character_count": len(input_text),
            "sentences": input_text.split('.'),
            "metadata": {
                "processed_by": "llama.cpp demo",
                "timestamp": "demo_mode"
            }
        }
        
        return json.dumps(result, indent=2)
        
    except Exception as e:
        return f"Error processing text: {str(e)}"

# Create Gradio interface
with gr.Blocks(title="Plain Text to JSON with llama.cpp") as demo:
    gr.Markdown("# Plain Text to JSON Converter")
    gr.Markdown("Convert plain text into structured JSON format using llama.cpp and Osmosis Structure model")
    
    with gr.Tab("Text to JSON"):
        with gr.Row():
            with gr.Column():
                input_text = gr.Textbox(
                    label="Input Text",
                    placeholder="Enter your text here...",
                    lines=5
                )
                
                with gr.Row():
                    max_tokens = gr.Slider(
                        minimum=50,
                        maximum=1000,
                        value=512,
                        label="Max Tokens"
                    )
                    temperature = gr.Slider(
                        minimum=0.1,
                        maximum=1.0,
                        value=0.7,
                        label="Temperature"
                    )
                
                convert_btn = gr.Button("Convert to JSON", variant="primary")
                demo_btn = gr.Button("Demo (No Model)", variant="secondary")
                
            with gr.Column():
                output_json = gr.Textbox(
                    label="Generated JSON",
                    lines=10,
                    interactive=False
                )
    
    with gr.Tab("Model Management"):
        load_btn = gr.Button("Load Model", variant="primary")
        model_status = gr.Textbox(
            label="Model Status",
            value="Model not loaded",
            interactive=False
        )
        
        gr.Markdown("""
        ### Instructions:
        1. Click "Load Model" to download and initialize the Osmosis Structure model
        2. Use "Demo (No Model)" for basic functionality without loading the AI model
        3. The Osmosis model is optimized for structured data extraction and JSON generation
        
        ### Notes:
        - Uses llama.cpp for efficient CPU inference
        - Osmosis Structure 0.6B model (~1.2GB) will be downloaded automatically
        - Model is specialized for converting unstructured text to structured formats
        - Adjust max_tokens and temperature for different output styles
        """)
    
    # Event handlers
    convert_btn.click(
        fn=text_to_json,
        inputs=[input_text, max_tokens, temperature],
        outputs=output_json
    )
    
    demo_btn.click(
        fn=demo_without_model,
        inputs=input_text,
        outputs=output_json
    )
    
    load_btn.click(
        fn=load_model,
        outputs=model_status
    )

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