from flask import Flask, request
import requests
import json
import os

from jwt_setup import get_installation_access_token


app = Flask(__name__)

def get_llm_output(query: str):
    LLM_output = requests.post('https://robinroy03-fury-engine.hf.space/api/google/generate', json={"query": query, "llm": "gemini-1.5-pro", "knn": 3, "stream": False})
    LLM_output = json.loads(LLM_output.text)
    LLM_output = LLM_output['response'] + '\n\n**References**\n' + LLM_output['references']

    return LLM_output


def get_root_comment_id(node_id: str) -> str:
    """
    Parameters
    ----------
    node_id: str
      The ID of the comment node.

    Returns
    -------
    root comment node_id: str
      The ID of the root comment node ID

    """
    token = get_installation_access_token(
        os.environ.get('CHATBOT_PUB_KEY'),
        os.environ.get('APP_ID'),
        os.environ.get('INSTALLATION_ID')
    )
    headers = {"Authorization": f"token {token}"}
    query = "query get_root_comment_ID ($id: ID!){node(id: $id) {... on DiscussionComment {replyTo {id}}}}"
    variables = {"id": node_id}

    json_return = requests.post('https://api.github.com/graphql', json={'query': query, 'variables': variables}, headers=headers)
    json_return = json.loads(json_return.text)

    return json_return['data']['node']['replyTo']['id']


def post_discussion_reply(discussion_id: str, replytoid: str, body: str):
    """
    Parameters
    ----------
    discussion_id: str
        ID of the discussion, available on the webhook payload
    replytoid: str
        ID of the root reply. This is not the ID of a threaded reply. Use get_root_comment_id() for the root ID.
    body: str
        Body of the reply

    """
    token = get_installation_access_token(
        os.environ.get('CHATBOT_PUB_KEY'),
        os.environ.get('APP_ID'),
        os.environ.get('INSTALLATION_ID')
    )

    headers = {"Authorization": f"token {token}"}
    query = "mutation AddDiscussionComment($discussionId: ID!, $body: String!, $replyToId: ID!) {addDiscussionComment(input: {discussionId: $discussionId, body: $body, replyToId: $replyToId}) {clientMutationId}}"
    variables = {"discussionId": discussion_id, "body": body, "replyToId": replytoid}

    requests.post('https://api.github.com/graphql', json={'query': query, 'variables': variables}, headers=headers)


def post_discussion_comment(discussion_id: str, body: str):
    token = get_installation_access_token(
        os.environ.get('CHATBOT_PUB_KEY'),
        os.environ.get('APP_ID'),
        os.environ.get('INSTALLATION_ID')
    )

    headers = {"Authorization": f"token {token}"}
    query = "mutation AddDiscussionComment($discussionId: ID!, $body: String!) {addDiscussionComment(input: {discussionId: $discussionId, body: $body}) {clientMutationId}}"
    variables = {"discussionId": discussion_id, "body": body}

    requests.post('https://api.github.com/graphql', json={'query': query, 'variables': variables}, headers=headers)


@app.route("/github", methods=['POST'])
def main():
    payload = request.get_json()
    URL = "https://api.github.com/graphql"

    if 'discussion' in payload and 'comment' not in payload:
        discussion_id = payload['discussion']['node_id']
        discussion_title = payload['discussion']['title']
        discussion_body = payload['discussion']['body']

        if ("[X] I want an AI-generated answer as a first response." in discussion_body) or ("[x] I want an AI-generated answer as a first response." in discussion_body):
            post_discussion_comment(discussion_id=discussion_id, body=get_llm_output(f"{discussion_title}\n{discussion_body}"))
        else:
            return {"status": False}    

        return {"status": True}
    
    elif 'discussion' in payload and 'comment' in payload:
        if payload['comment']['parent_id'] != None and payload['comment']['body'].startswith("!fury"):
            post_discussion_reply(
                payload['discussion']['node_id'],
                get_root_comment_id(payload['comment']['node_id']),
                get_llm_output(payload['comment']['body'][6:])    
            )
        elif payload['comment']['parent_id'] == None and payload['comment']['body'].startswith("!fury"):
            post_discussion_reply(
                payload['discussion']['node_id'],
                payload['comment']['node_id'],
                get_llm_output(payload['comment']['body'][6:])
            )

        return {"status": True}

    else:
        return {"status": False}