import os
import json

from flask import Flask, request
from flask_cors import CORS

import google.generativeai
google.generativeai.configure(
    api_key=os.environ.get("GEMINI_API_KEY"))

app = Flask(__name__)
CORS(app)


def output_sanitizer(prompt: str) -> dict:
    """
    return dictionary from llm output
    """
    begin = prompt.find("{")
    end = prompt.find("}")
    output = prompt[begin: end + 1]
    
    try:
        output = json.loads(output)
    except json.JSONDecodeError:
        return {}
    
    return output


def dict_style_changer(file: dict) -> dict: 
    """
    parses it in such a way so that mermaid won't crash due to semicolons and such
    """
    """
    {"1[a]" : ["2[b]", "3[c]"], "4[d]" : ["5[e]", "6[f]"]}
    """
    new = {}
    count = 1

    for k, v in file.items():
        new_key = f"{count}[{k}]"
        count += 1
        temp = []

        for i in v:
            new_value = f"{count}[{i}]"
            count += 1
            temp.append(new_value)
        
        new[new_key] = temp

    return new


def dict_to_mermaid(file: dict) -> str:
    """
    returns a mermaid.js syntax to be rendered for the roadmap
    """
    file = dict_style_changer(file)
    output = """
    flowchart TB\n\t
    """
    arrow = " --> "
    thick_arrow = " ==> "

    for k, v in file.items():
        output += f"subgraph {k}\n\t{arrow.join(v)}\n\tend\n\t"
    
    output += thick_arrow.join(list(file.keys()))

    return output


def get_completion(message: str, current_knowledge: str) -> str:
    """
    Returns a string with mermaid syntax.
    """
    sys_message = f"""
You are required to provide a learning roadmap for the given topic. The learner profile is provided. Use these data to give a very detailed learning roadmap. Use the output format shown below. Remember to optimize the roadmap as per the learners profile.

Return the output in JSON format.
Example:
{{
    "<Topic>": [
        "<sub topic 1>", 
        "<sub topic 2>"
    ],
    "<Topic>": [
        "<sub topic 1>", 
        "<sub topic 2>"
    ]
}}

Topic name: {message}
Learner profile: {current_knowledge}
"""
    llm_model = google.generativeai.GenerativeModel('gemini-1.5-pro')
    chat_completion = llm_model.generate_content(sys_message)
    print(chat_completion)
    chat_completion_text = chat_completion.to_dict()['candidates'][0]['content']['parts'][0]['text']
    chat_completion_text = chat_completion_text.replace("(", "")
    chat_completion_text = chat_completion_text.replace(")", "")

    print(dict_to_mermaid(output_sanitizer(chat_completion_text)))
    return dict_to_mermaid(output_sanitizer(chat_completion_text))



@app.route("/gemini_output", methods=['POST'])
def gemini_output():
    message = request.get_json()
    print(message)
    if message['prompt'] == "chess":
        return {'response': """flowchart TB

    subgraph 1[Fundamentals]
        2[Objective of the game] --> 3[Chessboard setup] --> 4[Piece movement King, Queen, Rook, Bishop, Knight, Pawn] --> 5[Special moves Castling, En Passant] --> 6[Check, Checkmate, and Stalemate]
        end
        subgraph 7[Basic Tactics]
        8[Piece value and development] --> 9[Controlling the center] --> 10[Forks] --> 11[Pins] --> 12[Skewers] --> 13[Discovered attacks] --> 14[Double attacks] --> 15[Traps]
        end
        subgraph 16[Openings]
        17[Principles of opening theory] --> 18[Common opening traps to avoid] --> 19[Learning a few standard openings e.g., Italian Game, Ruy Lopez, Sicilian Defense] --> 20[Understanding opening goals control of the center, piece development, king safety]
        end
        subgraph 21[Middlegame Strategy]
        22[Planning and evaluating positions] --> 23[Pawn structures and their weaknesses] --> 24[Identifying tactical opportunities] --> 25[Positional play improving piece placement, creating weaknesses] --> 26[Attacking and defending techniques]
        end
        subgraph 27[Endgames]
        28[Basic checkmates King and Queen, King and Rook] --> 29[Theoretical endgames King and Pawn vs King, Rook and Pawn vs Rook] --> 30[Key endgame principles opposition, triangulation, key squares] --> 31[Practice common endgame patterns]
        end
        subgraph 32[Beyond the Basics]
        33[Studying annotated master games] --> 34[Analyzing your own games] --> 35[Solving chess puzzles and tactics] --> 36[Participating in tournaments and online play] --> 37[Utilizing chess software and engines for analysis and training] --> 38[Exploring advanced concepts e.g., prophylaxis, positional sacrifice]
        end
        1[Fundamentals] ==> 7[Basic Tactics] ==> 16[Openings] ==> 21[Middlegame Strategy] ==> 27[Endgames] ==> 32[Beyond the Basics]"""}
    # return {"prompt": "hi " + message['prompt']}
    return {"response": get_completion(message['prompt'], message['current_knowledge'])}