from concurrent.futures import ThreadPoolExecutor, as_completed
import os
from openai import AzureOpenAI, OpenAI
import json
from story_schema import Story, StoryPage

client = AzureOpenAI(
  azure_endpoint = 'https://storymii-openai.openai.azure.com/',
  api_key="8u2PbSBg9U29jGIUJVYy67pWaKU2UYe9l1TnAZvtXKbTg62xr5KlJQQJ99ALACYeBjFXJ3w3AAABACOG2gJ2",
  api_version="2024-08-01-preview",
)

image_client = OpenAI(api_key = "sk-storymii-prod-images-R53ukBLUScAtCUoX2TAkT3BlbkFJO9JzB9NTYBUddahQj4Iw")

def generate_story(ageRange, character_count, style_story, story_theme, language, plot, main_character_name, main_character_attr, other_details, tone, diverse_learner) -> Story:
    try:
        if diverse_learner == "Dyslexic":
            system_prompt = system_dislecxic(ageRange, character_count)
            user_prompt = user_dislecxic(ageRange, character_count, style_story, story_theme, language, plot, main_character_name, main_character_attr, other_details, tone)
        elif diverse_learner == "ADHD":
            system_prompt = system_adhdprompt(ageRange, character_count)
            user_prompt = user_adhdprompt(ageRange, character_count, style_story, story_theme, language, plot, main_character_name, main_character_attr, other_details, tone)
        else:
            system_prompt = other_system_prompt(ageRange, character_count, diverse_learner)
            user_prompt = other_user_prompt(ageRange, character_count, style_story, story_theme, language, plot, main_character_name, main_character_attr, other_details, tone)

        response = client.chat.completions.create(
            model="gpt-4",
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_prompt}
            ]
        )

        story_content = response.choices[0].message.content
        story_data = json.loads(story_content)
        
        # Validate and structure the story data according to the schema
        story: Story = {
            'title': story_data['title'],
            'cover_image_prompt': story_data['cover_image_prompt'],
            'pages': [
                {
                    'content': page['content'],
                    'image_prompt': page['image_prompt']
                } for page in story_data['pages']
            ]
        }
        
        # Ensure exactly 4 pages
        if len(story['pages']) != 4:
            raise ValueError("Story must contain exactly 4 pages")
            
        return story

    except json.JSONDecodeError:
        raise ValueError("Invalid JSON response from story generation")
    except KeyError as e:
        raise ValueError(f"Missing required field in story response: {e}")
    except Exception as e:
        raise Exception(f"Error generating story: {str(e)}")

def system_dislecxic(ageRange, character_count):
    return f"""
    [Your existing system_dislecxic function content]
    """

def user_dislecxic(ageRange, character_count, style_story, story_theme, language, plot, main_character_name, main_character_attr, other_details, tone):
    return f"""
    [Your existing user_dislecxic function content]
    """

def system_adhdprompt(ageRange, character_count):
    return f"""
    [Your existing system_adhdprompt function content]
    """

def user_adhdprompt(ageRange, character_count, style_story, story_theme, language, plot, main_character_name, main_character_attr, other_details, tone):
    return f"""
    [Your existing user_adhdprompt function content]
    """

def other_system_prompt(ageRange, character_count, diverse_learner):
    return f"""
    [Your existing other_system_prompt function content]
    """

def other_user_prompt(ageRange, character_count, style_story, story_theme, language, plot, main_character_name, main_character_attr, other_details, tone):
    return f"""
    [Your existing other_user_prompt function content]
    """