import gradio as gr
import os
from huggingface_hub import HfApi
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline

# We'll use Mistral-7B as it's a good balance of quality and performance
MODEL_NAME = "mistralai/Mistral-7B-Instruct-v0.2"

# System prompt that incorporates the configuration
SYSTEM_PROMPT = """
You are Spectral Satya, a specialized AI assistant focused on crafting cinematic reality through expert prompt engineering for AI video generation.

Your primary purpose is to help users create highly realistic, professional-quality cinematic scene prompts for platforms including RunwayML, Pika, Kling, Haiper, Vidu, Veo, PixVerse, and other T2V/I2V models.

Core Principles:
- Realism Above All: Always prioritize photorealistic, cinematic quality over stylized or animated looks
- Specificity is King: Eliminate vagueness in all prompts; be precise about subjects, actions, environments, camera work, and lighting
- Show, Don't Tell: Use visual language that paints clear pictures rather than abstract concepts
- Defensive Prompting: Always include robust negative prompts to ward off unwanted styles and artifacts

Always structure prompts with these essential elements:
1. Style Declaration: Begin with "cinematic, photorealistic" or similar terms establishing the desired realism
2. Camera Specifications: Include shot type, angle, movement, and focus details
3. Subject Definition: Clearly define characters/subjects with specific visual attributes
4. Action Description: Detail precise movements and interactions
5. Environment Details: Describe setting with specific visual elements
6. Lighting & Atmosphere: Specify lighting conditions, mood, and atmospheric elements
7. Color & Grading: Include color palette or grading style information
8. Negative Prompts: Always provide comprehensive negative prompting to prevent unwanted styles

Your responses should follow the principles and techniques outlined in "The Prompt Engineer's Codex: Crafting Cinematic Reality with AI Video".
"""

# Initialize tokenizer and model
try:
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
    pipe = pipeline(
        "text-generation",
        model=MODEL_NAME,
        tokenizer=tokenizer,
        max_new_tokens=1024,
        temperature=0.7,
        top_p=0.95,
        repetition_penalty=1.15
    )
except Exception as e:
    print(f"Error loading model: {e}")
    # Fallback to a smaller model if the main one fails
    MODEL_NAME = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
    pipe = pipeline(
        "text-generation",
        model=MODEL_NAME,
        tokenizer=tokenizer,
        max_new_tokens=512,
        temperature=0.7,
        top_p=0.95,
        repetition_penalty=1.15
    )

def format_prompt(user_input, chat_history):
    """Format the prompt for the model with chat history."""
    messages = []
    
    # Add system prompt
    messages.append({"role": "system", "content": SYSTEM_PROMPT})
    
    # Add chat history
    for user_msg, assistant_msg in chat_history:
        messages.append({"role": "user", "content": user_msg})
        messages.append({"role": "assistant", "content": assistant_msg})
    
    # Add current user input
    messages.append({"role": "user", "content": user_input})
    
    # Format for the model
    formatted_prompt = ""
    for message in messages:
        if message["role"] == "system":
            formatted_prompt += f"<s>[INST] <<SYS>>\n{message['content']}\n<</SYS>>\n\n"
        elif message["role"] == "user":
            if formatted_prompt:
                formatted_prompt += f"{message['content']} [/INST]"
            else:
                formatted_prompt += f"<s>[INST] {message['content']} [/INST]"
        elif message["role"] == "assistant":
            formatted_prompt += f" {message['content']} </s><s>[INST] "
    
    return formatted_prompt

def generate_response(user_input, chat_history):
    """Generate a response using the model."""
    prompt = format_prompt(user_input, chat_history)
    
    try:
        response = pipe(prompt)[0]['generated_text']
        # Extract only the new content (the model's response)
        response = response.split('[/INST]')[-1].strip()
        if '</s>' in response:
            response = response.split('</s>')[0].strip()
        
        # Clean up any remaining tags
        response = response.replace('<s>', '').replace('</s>', '').replace('[INST]', '').replace('[/INST]', '')
        
        return response
    except Exception as e:
        return f"I apologize, but I encountered an error: {str(e)}. Please try again with a different query."

def spectral_satya_chat(user_input, chat_history):
    """Main chat function for the Gradio interface."""
    response = generate_response(user_input, chat_history)
    chat_history.append((user_input, response))
    return "", chat_history

def generate_cinematic_prompt(scene_description, subject_details, camera_preferences, lighting_mood, platform):
    """Generate a complete cinematic prompt based on user inputs."""
    prompt = f"""
    Please create a detailed cinematic prompt for the following scene:
    
    Scene Description: {scene_description}
    Subject Details: {subject_details}
    Camera Preferences: {camera_preferences}
    Lighting and Mood: {lighting_mood}
    Target Platform: {platform}
    
    I need both a positive prompt and a negative prompt that follows the principles from The Prompt Engineer's Codex.
    """
    
    chat_history = []
    response = generate_response(prompt, chat_history)
    
    return response

# Define example prompts
EXAMPLE_PROMPTS = """
# Example Cinematic Prompts

## Urban Character Scene