import gradio as gr
import os

# Try to import transformers, but provide fallback
try:
    from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
    TRANSFORMERS_AVAILABLE = True
except ImportError:
    TRANSFORMERS_AVAILABLE = False
    print("Transformers not available - running in demo mode")

def load_model():
    """Load the model if transformers is available"""
    if not TRANSFORMERS_AVAILABLE:
        return None
    
    try:
        model_id = "aliarsal1512/clarifai_java_code_commenter"
        print(f"Loading model: {model_id}")
        
        tokenizer = AutoTokenizer.from_pretrained(model_id)
        model = AutoModelForSeq2SeqLM.from_pretrained(model_id)
        
        pipe = pipeline(
            "text2text-generation",
            model=model,
            tokenizer=tokenizer,
            max_length=64,
            num_beams=1,
            do_sample=False
        )
        
        print("✅ Model loaded successfully!")
        return pipe
    except Exception as e:
        print(f"❌ Error loading model: {e}")
        return None

# Load model
pipe = load_model()

def generate_comment(code):
    """Generate comment for Java code"""
    if pipe is None:
        if not TRANSFORMERS_AVAILABLE:
            return "⚠️ Transformers library not installed. Please check requirements.txt"
        return "⚠️ Model failed to load. Check logs above."
    
    try:
        # Clean and prepare code
        code = code.strip()
        if not code:
            return "Please provide some Java code"
        
        # Generate comment
        result = pipe(code)
        comment = result[0]['generated_text']
        
        return comment
    except Exception as e:
        return f"❌ Error: {str(e)}"

# Create Gradio interface
demo = gr.Interface(
    fn=generate_comment,
    inputs=gr.Textbox(
        label="Java Code",
        lines=10,
        placeholder="Paste your Java code here...\nExample: public class Hello { public static void main(String[] args) { } }"
    ),
    outputs=gr.Textbox(
        label="Generated Comment",
        lines=4
    ),
    title="Java Code Comment Generator",
    description="Generates comments for Java code using a fine-tuned T5 model",
    examples=[
        ["public class Calculator {\n    public int add(int a, int b) {\n        return a + b;\n    }\n}"],
        ["public class Main {\n    public static void main(String[] args) {\n        System.out.println(\"Hello, World!\");\n    }\n}"]
    ],
    theme="soft"
)

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