# pnutty
# Copyright (C) 2024  Morgan McMillian

# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.

# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.

# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

import pnutpy
import sys
import os
import tomlkit
import cmd2
import websocket
import threading
import logging
import time
import json
import requests

from pathlib import Path
from colorama import Fore, Style

PNUT_CLIENT_ID = "qx87SsGu3NZINwAUxEk1bC1TICzupv5e"
PNUT_AUTH_URL = "https://pnut.io/oauth/authenticate"
PNUT_REDIRECT = "urn:ietf:wg:oauth:2.0:oob"
PNUT_SCOPES = "basic,stream,write_post,follow,presence,messages,files"
PNUT_WS_URL = "wss://stream.pnut.io/v1/user?access_token="
PNUT_API_URL = "https://api.pnut.io/v1"
PNUT_POST_STREAMS_UNIFIED = "/posts/streams/unified"
PNUT_POST_STREAMS_PERSONAL = "/posts/streams/me"
PNUT_POST_STREAMS_GLOBAL = "/posts/streams/global"
PNUT_POSTS_URL = "https://posts.pnut.io/"

REPLY_ONE = 1
REPLY_ALL = 2
REPLY_CC = 3

_shutdown = threading.Event()
_connected = threading.Event()

class PnutStream:

    def __init__(self, config):
        self.config = config
        self.token = config.get('token', '')
        self.show_links = config.get('show_links', False)
        self.timeline = config.get('timeline', 'global')
        self.pnut_api = pnutpy.API.build_api(access_token=self.token)
        self.ws = websocket.WebSocketApp(PNUT_WS_URL + self.token,
            on_message=self.on_message,
            on_error=self.on_error,
            on_close=self.on_close)
        self.ws.on_open = self.on_open

    def on_open(self, ws):
        logging.debug("--- starting stream connection with pnut.io ---")

        def run(*args):
            while not _shutdown.is_set():
                ws.send(".")
                time.sleep(2)
            logging.debug("*** terminating stream ***")

        t = threading.Thread(target=run)
        t.start()

    def on_close(self, ws, status_code, msg):
        logging.debug("--- closing stream connection with pnut.io ---")
        logging.debug(status_code)
        logging.debug(msg)

    def on_error(self, ws, error):
        logging.error("--- error on stream connection with pnut.io ---")
        logging.error(error)

    def on_message(self, ws, message):
        logging.debug("<on_message> ")
        msg = json.loads(message)

        if not _connected.isSet() and 'connection_id' in msg['meta']:
            connection_id = msg['meta']['connection_id']
            self.subscribe(connection_id)
            logging.debug("[connection_id]: " + connection_id)
            return

        if 'data' in msg:
            self.parse_data(msg)

    def subscribe(self, connection_id):
        # TODO: Replace with with pnutpy once it's updated
        if self.timeline == 'unified':
            url = PNUT_API_URL + PNUT_POST_STREAMS_UNIFIED
        elif self.timeline == 'personal':
            url = PNUT_API_URL + PNUT_POST_STREAMS_PERSONAL
        else:
            url = PNUT_API_URL + PNUT_POST_STREAMS_GLOBAL

        url += "?connection_id=" + connection_id
        headers = {'Authorization': 'Bearer ' + self.token}
        r = requests.get(url, headers=headers)
        if r.status_code == 200:
            _connected.set()
            self.parse_data(r.json())

    def parse_data(self, data):
        for item in reversed(data['data']):
            post = pnutpy.models.Post.from_response_data(item)
            if 'content' in post:
                self.print_post(post)

    def print_post(self, post):
        text = f"{Fore.YELLOW}{post.id}{Fore.RESET}: "
        if 'repost_of' in post:
            text += f"<{Style.BRIGHT}{Fore.CYAN}{post.user.username}{Style.RESET_ALL}> reposted >> "
            post = post.repost_of
        text += f"<{Style.BRIGHT}{Fore.CYAN}{post.user.username}{Style.RESET_ALL}>"
        text += f" {post.content.text} "
        text += f"{Fore.LIGHTBLACK_EX}[via {post.source.name}]{Style.RESET_ALL}"
        if self.show_links:
            text += f"\n{Fore.LIGHTBLACK_EX}{PNUT_POSTS_URL}{post.id}{Style.RESET_ALL}"
        print(text)

class NuttyShell(cmd2.Cmd):
    prompt = "pnutty% "

    def __init__(self, config):
        super().__init__()
        # self.allow_cli_args = False # defaults to True
        self.config = config
        self.stream = PnutStream(config)
        self.debug = config.get('debug', False)
        self.token = config.get('token')
        self.timeline = config.get('timeline', 'global')
        self.add_settable(
            cmd2.Settable('timeline', str,
                'Set timeline to [global, unified, or personal]',
                self, choices=['unified', 'personal', 'global'],
                onchange_cb=self._onchange_setting))
        self.feed = config.get('feed', True)
        self.add_settable(
            cmd2.Settable('feed', bool, 'Enable live feed',
                self, onchange_cb=self._onchange_setting))
        self.show_links = config.get('show_links', False)
        self.add_settable(
            cmd2.Settable('show_links', bool, 'Show link to post',
                self, onchange_cb=self._onchange_setting))
        self.register_preloop_hook(self.startup)
        self.register_postloop_hook(self.shutdown)

    def do_login(self, statement):
        '''Log in to pnut.io'''
        reply = "Visit the following URL to authorize with pnut.io and paste "
        reply += "the token here.\n\n"
        reply += f"{PNUT_AUTH_URL}?client_id={PNUT_CLIENT_ID}"
        reply += f"&redirect_uri={PNUT_REDIRECT}&scope={PNUT_SCOPES}"
        reply += f"&response_type=token\n"
        self.poutput(reply)
        token = input("TOKEN: ")
        self.check_auth(token)

    def do_logout(self, statement):
        '''Log out of pnut.io'''
        self.shutdown()
        url = PNUT_API_URL + '/token'
        headers = {'Authorization': 'Bearer ' + self.token}
        r = requests.delete(url, headers=headers)
        logging.debug(r.status_code)
        if r.status_code == 200:
            self.token = ""
            self.config['token'] = ""
            self.config['username'] = ""
            save_config(self.config)

    def do_again(self, statement):
        '''Replay your stream'''
        self.get_timeline()

    def do_global(self, statement):
        '''Switch timeline to global'''
        self.set_timeline('global')

    def do_unified(self, statement):
        '''Switch timeline to unified'''
        self.set_timeline('unified')

    def do_personal(self, statement):
        '''Switch timeline to personal'''
        self.set_timeline('personal')

    show_post_args = cmd2.Cmd2ArgumentParser()
    show_post_args.add_argument('post_id', help='ID of the post to show.')
    @cmd2.with_argparser(show_post_args)
    def do_show_post(self, args):
        pnut_api = pnutpy.API.build_api(access_token=self.token)

        try:
            post, meta = pnut_api.get_post(args.post_id)
            self.stream.print_post(post)

        except pnutpy.errors.PnutAuthAPIException:
            self.perror("Unable to access your account.")

    post_args = cmd2.Cmd2ArgumentParser()
    post_args.add_argument('text', type=str, nargs='+',
        help='text of message you wish to post.')
    @cmd2.with_argparser(post_args)
    def do_post(self, args):
        """Post to pnut.io"""
        text = ' '.join(args.text)
        self.send_post(text)

    reply_args = cmd2.Cmd2ArgumentParser()
    reply_args.add_argument('post_id', help='ID of the post to reply to.')
    reply_args.add_argument('text', nargs='+', help='text of your reply.')
    @cmd2.with_argparser(reply_args)
    def do_reply(self, args):
        '''Reply to a post'''
        text = ' '.join(args.text)
        self.send_post(text, args.post_id, REPLY_ONE)

    replyg_args = cmd2.Cmd2ArgumentParser()
    replyg_args.add_argument('post_id', help='ID of the post to reply to.')
    replyg_args.add_argument('text', nargs='+', help='text of your reply.')
    @cmd2.with_argparser(replyg_args)
    def do_replyg(self, args):
        '''Reply globally to a post'''
        text = ' '.join(args.text)
        self.send_post(text, args.post_id)

    replycc_args = cmd2.Cmd2ArgumentParser()
    replycc_args.add_argument('post_id', help='ID of the post to reply to.')
    replycc_args.add_argument('text', nargs='+', help='text of your reply.')
    @cmd2.with_argparser(replycc_args)
    def do_replycc(self, args):
        '''Reply to a post and cc others'''
        text = ' '.join(args.text)
        self.send_post(text, args.post_id, REPLY_CC)

    replyall_args = cmd2.Cmd2ArgumentParser()
    replyall_args.add_argument('post_id', help='ID of the post to reply to.')
    replyall_args.add_argument('text', nargs='+', help='text of your reply.')
    @cmd2.with_argparser(replyall_args)
    def do_replyall(self, args):
        '''Reply to all of a post'''
        text = ' '.join(args.text)
        self.send_post(text, args.post_id, REPLY_ALL)

    browse_args = cmd2.Cmd2ArgumentParser()
    browse_args.add_argument('post_id', nargs='+',
        help='ID of the post to open.')
    @cmd2.with_argparser(browse_args)
    def do_browse(self, args):
        '''Open a post in the web browser'''
        for post_id in args.post_id:
            url = PNUT_POSTS_URL + post_id
            self.poutput(url)

    bookmark_args = cmd2.Cmd2ArgumentParser()
    bookmark_args.add_argument('post_id', help='ID of the post to bookmark.')
    @cmd2.with_argparser(bookmark_args)
    def do_bookmark(self, args):
        '''Bookmark a post'''
        pnut_api = pnutpy.API.build_api(access_token=self.token)

        try:
            post, meta = pnut_api.bookmark_post(args.post_id)
            self.poutput(
                f"* bookmarked post {Fore.YELLOW}{post['id']}{Fore.RESET}")

        except pnutpy.errors.PnutAuthAPIException:
            self.perror("Unable to access your account.")

    unbookmark_args = cmd2.Cmd2ArgumentParser()
    unbookmark_args.add_argument('post_id',
        help='ID of the post to unbookmark.')
    @cmd2.with_argparser(unbookmark_args)
    def do_unbookmark(self, args):
        '''Remove a bookmark'''
        pnut_api = pnutpy.API.build_api(access_token=self.token)

        try:
            post, meta = pnut_api.bookmark_post(args.post_id)
            self.poutput(
                f"* removed bookmark post {Fore.YELLOW}{post['id']}{Fore.RESET}")

        except pnutpy.errors.PnutAuthAPIException:
            self.perror("Unable to access your account.")

    repost_args = cmd2.Cmd2ArgumentParser()
    repost_args.add_argument('post_id', help='ID of the post to repost.')
    @cmd2.with_argparser(repost_args)
    def do_repost(self, args):
        '''Repost a post'''
        pnut_api = pnutpy.API.build_api(access_token=self.token)

        try:
            post, meta = pnut_api.repost_post(args.post_id)
            self.poutput(
                f"* reposted {Fore.YELLOW}{post['id']}{Fore.RESET}")

        except pnutpy.errors.PnutAuthAPIException:
            self.perror("Unable to access your account.")

    unrepost_args = cmd2.Cmd2ArgumentParser()
    unrepost_args.add_argument('post_id', help='ID of the post to unrepost.')
    @cmd2.with_argparser(unrepost_args)
    def do_unrepost(self, args):
        '''Repost a post'''
        pnut_api = pnutpy.API.build_api(access_token=self.token)

        try:
            post, meta = pnut_api.unrepost_post(args.post_id)
            self.poutput(
                f"* unreposted {Fore.YELLOW}{post['id']}{Fore.RESET}")

        except pnutpy.errors.PnutAuthAPIException:
            self.perror("Unable to access your account.")

    def startup(self) -> None:
        if self.debug:
            logging.basicConfig(level=logging.DEBUG)
        if self.feed and len(self.token) > 0:
            if self.check_auth(self.token):
                t = threading.Thread(target=self.stream.ws.run_forever)
                t.start()

    def shutdown(self) -> None:
        _shutdown.set()
        self.stream.ws.close()
        time.sleep(2)
        _shutdown.clear()
        _connected.clear()

    def check_auth(self, token):
        pnut_api = pnutpy.API.build_api(access_token=token)

        try:
            pnut_user, meta = pnut_api.get_user("me")
            self.config['username'] = pnut_user.username
            self.config['token'] = token
            save_config(self.config)

            reply = f"You have been successfully authenticated as "
            reply += f"{self.config['username']}."
            self.poutput(reply)
            return True

        except pnutpy.errors.PnutAuthAPIException:
            self.perror("Unable to access your account.")

        return False

    def get_timeline(self):
        pnut_api = pnutpy.API.build_api(access_token=self.token)

        try:
            if self.timeline == 'unified':
                data, meta = pnut_api.users_post_streams_unified()
            elif self.timeline == 'personal':
                data, meta = pnut_api.users_post_streams_me()
            else:
                data, meta = pnut_api.posts_streams_global()

            for post in reversed(data):
                if 'content' in post:
                    self.stream.print_post(post)

        except pnutpy.errors.PnutAuthAPIException:
            self.perror("Unable to access your account.")

    def set_timeline(self, value):
        self.shutdown()
        self.timeline = value
        self.stream.timeline = value
        self.startup()
        self.config['timeline'] = value
        save_config(self.config)

    def send_post(self, text, reply_to=None, mention=None):
        pnut_api = pnutpy.API.build_api(access_token=self.token)
        try:

            if mention is not None:
                orig_msg, meta = pnut_api.get_post(reply_to)
                highlight = "@" + orig_msg.user.username

                if mention == REPLY_ONE:
                    text = highlight + " " + text

                elif mention in [REPLY_ALL, REPLY_CC]:

                    cc_list = []
                    if 'content' in orig_msg:
                        if 'mentions' in orig_msg.content.entities:
                            mentions = orig_msg.content.entities.mentions
                        else:
                            mentions = []

                        for m in mentions:
                            if (m.text != self.config['username'] and
                                    m.text != highlight):
                                cc_list.append("@" + m.text)
                    if len(cc_list) > 0:
                        cc = " ".join(cc_list)

                        if mention == REPLY_CC:
                            text = highlight + " " + text + " /" + cc

                        elif mention == REPLY_ALL:
                            text = highlight + " " + cc + " " + text

            post, meta = pnut_api.create_post(
                data={'reply_to': reply_to, 'text': text})
            self.poutput(
                f"* created post {Fore.YELLOW}{post['id']}{Fore.RESET}")

        except pnutpy.errors.PnutAuthAPIException:
            self.perror("Unable to access your account.")

    def _onchange_setting(self, setting, old, new):
        if setting == 'feed' and new is True:
            self.startup()

        elif setting == 'feed' and new is False:
            self.shutdown()

        elif setting == 'timeline':
            self.shutdown()
            self.stream.timeline = new
            self.startup()

        elif setting == 'show_links':
            self.stream.show_links = new

        self.config[setting] = new
        save_config(self.config)

def save_config(config):
    config_file = get_config_file()
    with open(config_file, "w") as cf:
        tomlkit.dump(config, cf)

def get_config_file():
    XDG_CONFIG_HOME = os.environ.get("XDG_CONFIG_HOME")

    if bool(XDG_CONFIG_HOME):
        config_file = os.path.join(XDG_CONFIG_HOME, "pnutty.toml")
    else:
        config_file = Path.home().joinpath(".pnutty.toml")

    return Path(config_file)

def main():
    parser = cmd2.Cmd2ArgumentParser()
    parser.add_argument('-d', '--debug', action='store_const',
        dest='loglevel', const=logging.DEBUG, default=logging.INFO,
        help="enable debug log output")
    args = parser.parse_args()

    logging.basicConfig(level=args.loglevel)

    config_file = get_config_file()
    if config_file.exists():
        with open(config_file, "rb") as cf:
            config = tomlkit.load(cf)
    else:
        config = tomlkit.document()

    if 'token' not in config:
        config['token'] = ""

    app = NuttyShell(config)
    sys.exit(app.cmdloop())

if __name__ == '__main__':
    main()
