import gradio as gr
import torch
from models.text_gen import TextGenerator
from models.summarizer import TextSummarizer
from models.image_gen import ImageGenerator
from models.audio_gen import AudioGenerator
from models.code_gen import CodeGenerator

# Debug GPU availability
print(f"CUDA available: {torch.cuda.is_available()}")
if torch.cuda.is_available():
    print(f"CUDA device count: {torch.cuda.device_count()}")
    print(f"Current CUDA device: {torch.cuda.current_device()}")
    print(f"CUDA device name: {torch.cuda.get_device_name(0)}")
    print(f"CUDA device properties: {torch.cuda.get_device_properties(0)}")

# Initialize the models
text_generator = TextGenerator()
text_summarizer = TextSummarizer()
image_generator = ImageGenerator()
audio_generator = AudioGenerator()
code_generator = CodeGenerator()

def generate_text(prompt, max_length, temperature, top_p):
    try:
        generated_text = text_generator.generate_text(
            prompt=prompt,
            max_length=max_length,
            temperature=temperature,
            top_p=top_p
        )
        return generated_text
    except Exception as e:
        return f"Error generating text: {str(e)}"

def summarize_text(text, max_length, min_length):
    try:
        summary = text_summarizer.summarize(
            text=text,
            max_length=max_length,
            min_length=min_length
        )
        return summary
    except Exception as e:
        return f"Error generating summary: {str(e)}"

def generate_image(prompt, num_steps, guidance_scale):
    try:
        image = image_generator.generate_image(
            prompt=prompt,
            num_inference_steps=num_steps,
            guidance_scale=guidance_scale
        )
        return image
    except Exception as e:
        return f"Error generating image: {str(e)}"

def generate_audio(text):
    try:
        audio, sample_rate = audio_generator.generate_audio(
            text=text
        )
        return (sample_rate, audio)
    except Exception as e:
        return f"Error generating audio: {str(e)}"

def generate_code(prompt, max_length, temperature, top_p):
    try:
        code = code_generator.generate_code(
            prompt=prompt,
            max_length=max_length,
            temperature=temperature,
            top_p=top_p
        )
        return code
    except Exception as e:
        return f"Error generating code: {str(e)}"

# Create the Gradio interface
with gr.Blocks(title="GenAI Content Studio") as app:
    gr.Markdown("# 🎨 GenAI Content Studio")
    gr.Markdown("### Free and Open-Source AI Content Generation")
    
    with gr.Tabs():
        with gr.TabItem("Text Generation"):
            with gr.Row():
                with gr.Column():
                    prompt = gr.Textbox(
                        label="Enter your prompt",
                        placeholder="Type your text here...",
                        lines=3
                    )
                    
                    with gr.Row():
                        max_length = gr.Slider(
                            minimum=50,
                            maximum=500,
                            value=100,
                            step=50,
                            label="Max Length"
                        )
                        temperature = gr.Slider(
                            minimum=0.1,
                            maximum=1.0,
                            value=0.7,
                            step=0.1,
                            label="Temperature"
                        )
                        top_p = gr.Slider(
                            minimum=0.1,
                            maximum=1.0,
                            value=0.9,
                            step=0.1,
                            label="Top P"
                        )
                    
                    generate_btn = gr.Button("Generate Text")
                
                with gr.Column():
                    output = gr.Textbox(
                        label="Generated Text",
                        lines=10,
                        interactive=False
                    )
            
            generate_btn.click(
                fn=generate_text,
                inputs=[prompt, max_length, temperature, top_p],
                outputs=output
            )

        with gr.TabItem("Text Summarization"):
            with gr.Row():
                with gr.Column():
                    text_input = gr.Textbox(
                        label="Enter text to summarize",
                        placeholder="Paste your text here...",
                        lines=10
                    )
                    
                    with gr.Row():
                        max_summary_length = gr.Slider(
                            minimum=50,
                            maximum=200,
                            value=130,
                            step=10,
                            label="Max Summary Length"
                        )
                        min_summary_length = gr.Slider(
                            minimum=10,
                            maximum=100,
                            value=30,
                            step=10,
                            label="Min Summary Length"
                        )
                    
                    summarize_btn = gr.Button("Summarize Text")
                
                with gr.Column():
                    summary_output = gr.Textbox(
                        label="Generated Summary",
                        lines=5,
                        interactive=False
                    )
            
            summarize_btn.click(
                fn=summarize_text,
                inputs=[text_input, max_summary_length, min_summary_length],
                outputs=summary_output
            )
        
        with gr.TabItem("Image Generation"):
            with gr.Row():
                with gr.Column():
                    image_prompt = gr.Textbox(
                        label="Enter your image prompt",
                        placeholder="Describe the image you want to generate...",
                        lines=3
                    )
                    
                    with gr.Row():
                        num_steps = gr.Slider(
                            minimum=20,
                            maximum=100,
                            value=50,
                            step=10,
                            label="Number of Steps"
                        )
                        guidance_scale = gr.Slider(
                            minimum=1.0,
                            maximum=20.0,
                            value=7.5,
                            step=0.5,
                            label="Guidance Scale"
                        )
                    
                    generate_image_btn = gr.Button("Generate Image")
                
                with gr.Column():
                    image_output = gr.Image(
                        label="Generated Image",
                        type="pil"
                    )
            
            generate_image_btn.click(
                fn=generate_image,
                inputs=[image_prompt, num_steps, guidance_scale],
                outputs=image_output
            )

        with gr.TabItem("Audio Generation"):
            with gr.Row():
                with gr.Column():
                    audio_text = gr.Textbox(
                        label="Enter text to convert to speech",
                        placeholder="Type what you want to hear...",
                        lines=3
                    )
                    
                    generate_audio_btn = gr.Button("Generate Audio")
                
                with gr.Column():
                    audio_output = gr.Audio(
                        label="Generated Audio",
                        type="numpy"
                    )
            
            generate_audio_btn.click(
                fn=generate_audio,
                inputs=[audio_text],
                outputs=audio_output
            )
        
        with gr.TabItem("Code Generation"):
            with gr.Row():
                with gr.Column():
                    code_prompt = gr.Textbox(
                        label="Enter your code prompt",
                        placeholder="Describe the code you want to generate...",
                        lines=3
                    )
                    
                    with gr.Row():
                        max_length = gr.Slider(
                            minimum=50,
                            maximum=500,
                            value=100,
                            step=50,
                            label="Max Length"
                        )
                        temperature = gr.Slider(
                            minimum=0.1,
                            maximum=1.0,
                            value=0.7,
                            step=0.1,
                            label="Temperature"
                        )
                        top_p = gr.Slider(
                            minimum=0.1,
                            maximum=1.0,
                            value=0.9,
                            step=0.1,
                            label="Top P"
                        )
                    
                    generate_code_btn = gr.Button("Generate Code")
                
                with gr.Column():
                    code_output = gr.Code(
                        label="Generated Code",
                        language="python"
                    )
            
            generate_code_btn.click(
                fn=generate_code,
                inputs=[code_prompt, max_length, temperature, top_p],
                outputs=code_output
            )

if __name__ == "__main__":
    app.launch()