import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
from huggingface_hub import login
import os
import logging
from datetime import datetime

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Environment variables
HF_TOKEN = os.getenv("HUGGING_FACE_TOKEN")
MODEL_NAME = os.getenv("MODEL_NAME", "google/gemma-2b-it")

class CodeReviewer:
    def __init__(self):
        self.model = None
        self.tokenizer = None
        self.device = "cpu"
        self.initialize_model()
        
    def initialize_model(self):
        """Initialize the model and tokenizer."""
        try:
            if HF_TOKEN:
                login(token=HF_TOKEN)
            
            logger.info("Loading tokenizer...")
            self.tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
            
            logger.info("Loading model...")
            self.model = AutoModelForCausalLM.from_pretrained(
                MODEL_NAME,
                device_map={"": self.device},
                torch_dtype=torch.float32,
                low_cpu_mem_usage=True
            )
            logger.info("Model loaded successfully")
        except Exception as e:
            logger.error(f"Error initializing model: {e}")
            raise

    def create_review_prompt(self, code: str, language: str) -> str:
        """Create a structured prompt for code review."""
        return f"""Review this {language} code. List specific points in these sections:
Issues:
Improvements:
Best Practices:
Security:

Code:
```{language}
{code}
```"""

    def review_code(self, code: str, language: str) -> str:
        """Perform code review using the model."""
        try:
            prompt = self.create_review_prompt(code, language)
            
            inputs = self.tokenizer(
                prompt,
                return_tensors="pt",
                truncation=True,
                max_length=512,
                padding=True
            )
            
            with torch.no_grad():
                outputs = self.model.generate(
                    **inputs,
                    max_new_tokens=512,
                    do_sample=True,
                    temperature=0.7,
                    top_p=0.95,
                    num_beams=1,
                    early_stopping=True
                )
            
            response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
            return response[len(prompt):].strip()
            
        except Exception as e:
            logger.error(f"Error during code review: {e}")
            return f"Error performing code review: {str(e)}"

# Initialize the reviewer
reviewer = CodeReviewer()

def review_code_interface(code: str, language: str) -> str:
    """Gradio interface function for code review."""
    if not code.strip():
        return "Please enter some code to review."
    
    try:
        result = reviewer.review_code(code, language)
        return result
    except Exception as e:
        return f"Error: {str(e)}"

# Create Gradio interface
iface = gr.Interface(
    fn=review_code_interface,
    inputs=[
        gr.Textbox(
            lines=10,
            placeholder="Enter your code here...",
            label="Code"
        ),
        gr.Dropdown(
            choices=["python", "javascript", "java", "cpp", "typescript", "go", "rust"],
            value="python",
            label="Language"
        )
    ],
    outputs=gr.Textbox(
        label="Review Results",
        lines=10
    ),
    title="Code Review Assistant",
    description="An automated code review system powered by Gemma-2b that provides intelligent code analysis and suggestions for improvements.",
    examples=[
        ["""def add_numbers(a, b):
    return a + b""", "python"],
        ["""function calculateSum(numbers) {
    let sum = 0;
    for(let i = 0; i < numbers.length; i++) {
        sum += numbers[i];
    }
    return sum;
}""", "javascript"]
    ],
    theme=gr.themes.Soft()
)

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