from datetime import date
from typing import cast

import openai
import streamlit as st

import gcal
import google_oauth

st.title(":calendar: Creative Calendar")
st.subheader("Illustrate a month from your Google calendar")


client: openai.OpenAI | None = None
month_date: date | None = None
name: str | None = None
gender: str | None = None
genre: str | None = None
submitted = False

creds = google_oauth.load_creds()

# no Google credentials? we need to authenticate
if creds is None:
    flow = google_oauth.get_flow()
    url = google_oauth.get_auth_url(flow)
    st.link_button("Authenticate with Google", url=url)

    auth_code = st.text_input("Google auth code", type="password")
    if auth_code:
        creds = google_oauth.get_creds(flow, auth_code)
        st.rerun()

# we have Google credentials, make sure we have the other things
else:
    st.checkbox("Authenticated with Google", value=True, disabled=True)
    openapi_key = st.text_input("OpenAI API key", type="password")
    if openapi_key:
        client = openai.OpenAI(api_key=openapi_key)

        with st.form("my_form"):
            name = st.text_input("Your name")
            gender = st.radio("Your gender", ["Male", "Female"]) or "Male"
            genre = st.text_input("Illustration genre", value="fantasy")
            month_date = cast(date, st.date_input("Select month"))
            submitted = st.form_submit_button("Submit")


def compose_story_prompt(
    *,
    name: str,
    gender: str,
    events: list[gcal.Event],
    genre: str,
) -> str:
    prompt_events = []
    for ev in events:
        dt = ev.start.strftime("%B %d, %Y")
        line = f"{dt}: {ev.title}"
        if ev.description:
            line += f" ({ev.description})"
        prompt_events.append(line)

    event_str = "\n".join(["-" + ev for ev in prompt_events])
    prompt = f"""Consider the following events from my Google calendar:

    {event_str}

    Write a description for a short {genre} story that captures the major events
    from my calendar. The description should cover only the major themes from my
    calendar. Do not write the story itself, just a description of the story. Do
    not include any references to a calendar or specific dates, just the major
    themes from my calendar.
    
    The main character is a {gender} named {name}.
    """
    return prompt


def compose_picture_prompt(story: str, genre: str) -> str:
    prompt = f"""For the following short {genre} story description, create a single
    {genre}-style image that captures the major events in the story. Do not
    include any text in the illustration.
    -----

    {story}
    """
    return prompt


@st.cache_data(show_spinner="Writing the story...")
def write_story(story_prompt):
    completion = client.chat.completions.create(
        model="gpt-4",
        messages=[
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": story_prompt},
                ],
            }
        ],
    )
    content = completion.choices[0].message.content
    return content


@st.cache_data(show_spinner="Illustrating the cover...")
def draw_picture(picture_prompt: str) -> str:
    response = client.images.generate(
        model="dall-e-3",
        prompt=picture_prompt,
        size="1024x1024",
        quality="standard",
        n=1,
    )
    image_url = response.data[0].url
    return cast(str, image_url)


if client and creds and submitted:
    cal_service = gcal.build_calendar_api(creds)
    calendars = gcal.get_calendars(cal_service)

    month_start, month_end = gcal.get_start_and_end(cast(date, month_date))
    events = gcal.get_events(cal_service, month_start, month_end, calendars)

    st.write(
        f":white_check_mark: Fetched {len(events)} events from "
        f"{len(calendars)} calendars"
    )

    story_prompt = compose_story_prompt(
        name=cast(str, name),
        gender=cast(str, gender),
        events=events,
        genre=cast(str, genre),
    )
    story = write_story(story_prompt)

    with st.expander("See the story"):
        st.write(story)

    st.write(
        "Tip: The story and image quality depends on the quality"
        " of your calendar events."
    )
    picture_prompt = compose_picture_prompt(story, genre)
    image_url = draw_picture(picture_prompt)
    st.image(image_url)