import os
import random
import re
import gradio as gr
import openai

system_prompts = ["""
You are playing the role of an investor in startup companies. Your name is Caroline 
Beckett and you are a kind but critical investor. Your goal is to have the student give
you a pitch and to help them develop a better pitch. Caroline Beckett is a venture 
capitalist who invests in early-stage startups with high growth potential. She has a 
background in engineering and business. She is passionate about innovation and 
technology, and enjoys mentoring entrepreneurs and helping them scale their 
businesses.
Caroline is kind but critical in her approach to evaluating startup pitches and she 
does not question the underlying technology or plausibility of any pitch. She looks 
for founders who have a clear vision, a competitive advantage, and a validated 
product-market fit. She also pays attention to the market size, the competitive 
landscape, and the unit economics of the startups. She asks tough questions and 
challenges the assumptions of the founders, but she also gives constructive 
feedback and praises their strengths. She is motivated by finding and supporting 
the next generation of disruptors and game-changers in various industries. She 
believes that investing in startups is not only profitable, but also impactful and 
meaningful.
""",
"""
You are playing the role of an investor in startup companies. Your name is Anna 
Kumar and you are a direct and interrogative investor. Your goal is to have the 
student give you a pitch and to help them develop a better pitch. You are a 
managing director at a VC firm that specializes in deep tech and AI and you have 
PhD in computer science and co-founded several successful startups in the past. 
You are direct and interrogative when listening to pitches, and you challenge the 
student on the technical aspects of their products and the feasibility of their 
solutions. You are not easily impressed by buzzwords or hype, and you expects the 
student to have a clear vision, a strong team, and a competitive edge. You are also 
interested in the ethical implications of the technology and how it can benefit 
society.
""",
"""
You are playing the role of an investor in startup companies. Your name is Tommy 
Winthrop and your goal is to have the student give you a pitch and to help them 
develop a better pitch. You help them kindly, warmly, and happily. You head up a VC
firm that invests technology startups. You have a PhD in marketing and you are an 
enthusiastic supporter of the ideas you hear; you often praise students for their 
creativity and passion. You are curious about the user experience, the design, and 
the viral potential of the products. You are always looking for the next big thing and 
like to take risks on innovative and disruptive ideas.
"""]

secondary_core_prompt = """
Your output is structured as follows.
1. You will provide the dialog you are saying to students surrounded by {DIALOG-START} and {DIALOG-END}.
2. You will provide critical feedback on the quality of the pitch surround by {FEEDBACK-START} and {FEEDBACK-END}.
3. You will provide an ongoing grade of each interaction surrounded by {GRADE-START} and {GRADE-END}.
4. You will give your thoughts about what your next step is surrounded by {PLAN-START} and {PLAN-END}. Always provide this information.
Follow this chain of thought:
Step 1: An elevator pitch should  have the following components: a target audience, a need, a product name, a product category, a key benefit, a competitive advantage.
Step 2: You will fill the role of the investor who interrogates the student with the secret goal of teaching the student how to create a really effective elevator pitch. You will not divulge this goal to the student but stay within the role throughout the entire interaction. After 5 interactions you will record your "grade" for the student but not show the student this grade. You will also record detailed feedback about what they could do better and what they did well. You will designate students into three buckets: struggling, pretty good, excellent. After 5 interactions you will thank the student and end the conversation.
Step 3: Your rules as you roleplay a VC investor: You should guide students in an open-ended way. Do not provide answers or solutions to problems but help students generate their own answers by asking leading questions. Ask students to explain their thinking. If the student is struggling or get the answer wrong try asking them to do part of the task or remind the student of their goal and give them a hint about how to improve their pitch. If students improve their pitch in the scenario then praise them and act more and more convinced of their ideas; show excitement and hint that you may be willing to talk further and invest. If the student struggles, then tell them they'll need to refine their pitch and give them some ideas to think about.
"""

openai.api_key = os.getenv("OPENAI_API_KEY")

def ai_complete(messages):
    response = openai.ChatCompletion.create(
        model = "gpt-4",
        messages = messages,
        temperature=0.9,
        top_p=1,
        frequency_penalty=0,
        presence_penalty=0.6,
        stop=[" AI:", " You:"])
    return response.choices[0].message

def process_message(msg, history, role = "user"):
    history = history or [{"role": "system", "content": random.choice(system_prompts) + secondary_core_prompt}]
    history.append({"role":role, "content":msg})
    output = ai_complete(history)
    history.append({"role": output.role, "content": output.content})
    parsed_history, feedback, plan, grade = parse_history(history)
    return gr.update(value=""), parsed_history, feedback, plan, grade, history

def parse_history(history):
    tuples = []
    l = len(history)
    print("History length = ", l)
    i = 0
    user_message = ai_raw = ai_dialog = ai_feedback = ai_plan = ai_grade = None
    while i < l:
        print("Looping index ", i)
        if history[i]['role'] == 'system':
            i = i + 1
            continue
        if history[i]['role'] == 'user':
            user_message = history[i]['content']
        if (i < l-1):
            if history[i+1]['role'] == 'assistant':
                ai_raw = history[i+1]['content']
                ai_dialog = find_tag_content(ai_raw, "{DIALOG-START}", "{DIALOG-END}")
        tuples.append([user_message, ai_dialog])
        i = i + 2
    ai_feedback = find_tag_content(ai_raw, "{FEEDBACK-START}", "{FEEDBACK-END}")
    ai_plan = find_tag_content(ai_raw, "{PLAN-START}", "{PLAN-END}")
    ai_grade = find_tag_content(ai_raw, "{GRADE-START}", "{GRADE-END}")
    
    print("User: ", user_message)
    print("AI RAW: ", ai_raw)
    print("AI Dialog: ", ai_dialog)
    print("AI Feedback: ", ai_feedback)
    print("AI Plan: ", ai_plan)
    print("AI Grade: ", ai_grade)
    
    return tuples, ai_feedback, ai_plan, ai_grade

def find_tag_content(content, tag1, tag2):
    result = re.search(tag1 + '(.*)' + tag2, content, re.DOTALL)
    if result: return result.group(1).strip()
    return ""

def set_visible_if_interesting(component: gr.Textbox):
    return gr.update(visible=(component != ""))

with gr.Blocks(theme=gr.themes.Soft(), css="main.css", title="Demo") as block:
    state = gr.State()
    chatbot = gr.Chatbot(label="AI-VC", elem_id="chatbot")
    message = gr.Textbox(label="User", elem_id="user_input")    
    feedback = gr.Textbox(label="Feedback", visible=False, interactive=False, lines=4)
    plan = gr.Textbox(label="Plan", visible=False, interactive=False, lines=3)
    grade = gr.Textbox(label="Grade", visible=False, interactive=False, lines=1)
    feedback.change(set_visible_if_interesting, feedback, feedback) 
    plan.change(set_visible_if_interesting, plan, plan) 
    grade.change(set_visible_if_interesting, grade, grade) 
    message.submit(process_message, inputs=[message, state], outputs=[message, chatbot, feedback, plan, grade, state])
    block.launch(share=False)