from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
from langchain_core.messages import AIMessage

MODEL_REPO = "Rahul-8799/software_engineer_mellum"

tokenizer = AutoTokenizer.from_pretrained(MODEL_REPO, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_REPO,
    torch_dtype=torch.float16,
    device_map="auto"
)

def run(state: dict) -> dict:
    """Software Engineer generates clean, modern UI code using best practices"""
    messages = state["messages"]
    prompt = messages[-1].content
    
    # Enhance the prompt with UI implementation guidelines
    enhanced_prompt = f"""
    Objective

Generate modern, responsive, and accessible UI code that is visually appealing and adheres to current frontend development best practices.

1. Styling Framework: Tailwind CSS
	•	Use Tailwind CSS utility classes for styling all elements.
	•	Apply spacing, typography, sizing, and layout using Tailwind classes.
	•	Follow a mobile-first design approach.
	•	Use Tailwind’s built-in responsive breakpoints (sm, md, lg, xl, 2xl) to adapt layouts for different screen sizes.

2. Layout Techniques
	•	Use CSS Grid for complex, multi-column or two-dimensional layouts.
	•	Use Flexbox for flexible alignment of components like navigation bars, cards, buttons, and modals.
	•	Maintain consistent spacing with utility classes such as gap, space-x, space-y, p-*, and m-*.

3. Semantic HTML
	•	Use semantic HTML tags appropriately: <header>, <nav>, <main>, <section>, <article>, <footer>, etc.
	•	Avoid unnecessary <div> elements to prevent cluttered and unstructured markup.
	•	Ensure proper nesting and hierarchy of elements.

4. Accessibility
	•	Add ARIA labels, role attributes, and alt text where needed for screen reader support.
	•	Ensure keyboard accessibility with tabindex, proper focus states, and interactive elements being navigable.
	•	Use <label> elements properly linked to form fields via the for attribute.

5. Responsive Design
	•	Use Tailwind’s responsive utilities to adjust layouts across various screen sizes.
	•	Design components to be fully usable on both desktop and mobile devices.
	•	Use collapsible or toggleable UI patterns (e.g., hamburger menus) for smaller viewports.

6. Theming and Styling Consistency
	•	Define and use CSS variables (--primary-color, --font-family, etc.) for theme consistency across components.
	•	Maintain a clear visual hierarchy with consistent font sizes, weights, and colors.
	•	Customize Tailwind’s theme configuration if needed for project-specific design tokens.

7. JavaScript and Interactivity
	•	Add interactivity using plain JavaScript, Alpine.js, or React if specified.
	•	Implement common UI components such as modals, dropdowns, tooltips, accordions with appropriate open/close behavior.
	•	Provide user feedback through form validations, dynamic updates, and transitions.

8. Loading and Error States
	•	Implement loading states using spinners, skeleton screens, or placeholders while data is being fetched or actions are processing.
	•	Show error states using alerts, banners, or toast messages when applicable.
	•	Use conditional rendering or state flags to handle visibility and transitions between states.

9. Component Structure and Reusability
	•	Break down the UI into modular, reusable components (e.g., Button, Card, Modal, Form).
	•	Each component should:
	•	Be self-contained with a clear purpose.
	•	Accept inputs or props when necessary.
	•	Maintain responsive and accessible markup by default.

10. Code Quality Standards
	•	Write clean, readable, and maintainable code.
	•	Remove unused classes, scripts, or markup.
	•	Follow consistent naming conventions and indentation rules.
	•	Add comments only when necessary for clarity.


    Original requirements: {prompt}
    """
    
    input_ids = tokenizer(enhanced_prompt, return_tensors="pt").input_ids.to(model.device)
    output_ids = model.generate(input_ids, max_new_tokens=3000)
    output = tokenizer.decode(output_ids[0], skip_special_tokens=True)

    return {
        "messages": [AIMessage(content=output)],
        "chat_log": state["chat_log"] + [{"role": "Software Engineer", "content": output}],
        "dev_output": output,
    }