import os
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
from langchain_community.llms import HuggingFacePipeline
from langchain.chains import LLMChain
from prompts import prompt_template
from visual_generation import generate_design_image, create_figma_template

# Fetch environment variables
HUGGINGFACE_API_TOKEN = os.getenv("HUGGINGFACE_API_TOKEN")
FIGMA_API_KEY = os.getenv("FIGMA_API_KEY")

# Model name for GPT-Neo 1.3B
model_name = "EleutherAI/gpt-neo-1.3B"

# Initialize the tokenizer and model with the Hugging Face API token
tokenizer = AutoTokenizer.from_pretrained(model_name, token=HUGGINGFACE_API_TOKEN)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    token=HUGGINGFACE_API_TOKEN,
    device_map=None,  # Ensures CPU-only usage
    torch_dtype="float32"  # Use 32-bit precision for compatibility
)

# Create a Hugging Face pipeline for text generation
pipe = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    max_new_tokens=50,  # Limit the number of generated tokens
    device=-1  # device=-1 ensures CPU usage
)

# Wrap the pipeline with LangChain's HuggingFacePipeline
llm = HuggingFacePipeline(pipeline=pipe, verbose=True)

def generate_response(task, style, requirements):
    """
    Generate both textual and visual outputs based on the user's input.
    """
    # Generate textual suggestions using LangChain
    chain = LLMChain(llm=llm, prompt=prompt_template)
    text_response = chain.run(task=task, style=style, requirements=requirements)
    
    # Generate visual design using Stable Diffusion
    visual_prompt = f"Create a {style} design for {task} with {requirements}."
    visual_image = generate_design_image(visual_prompt)
    
    # Generate Figma template link
    figma_link = create_figma_template(f"{style} design for {task}")
    
    return {
        "text": text_response,
        "visual_template": visual_image,
        "figma_template": figma_link,
    }

# Test the chatbot locally
if __name__ == "__main__":
    task = "landing page"
    style = "modern minimalist"
    requirements = "responsive layout and clear CTAs"
    response = generate_response(task, style, requirements)
    print("Textual Suggestions:", response["text"])
    print("Visual Template URL:", response["visual_template"])
    print("Figma Template URL:", response["figma_template"])