#!/usr/bin/env python3

# to override print <= can be a big problem with exceptions
from __future__ import print_function  # must be 1st
import pandas as pd
import builtins
import datetime as dt
import os
import sys
import time
import json

from console import bg, fg, fx
from fire import Fire
from openai import OpenAI
from prompt_toolkit import PromptSession, prompt
from prompt_toolkit.history import FileHistory

import re


from cmd_ai import config # , key_enter, topbar, unitname
from cmd_ai import texts
from cmd_ai.api_key import get_api_key
from cmd_ai.g_askme import g_askme
from cmd_ai.g_askdalle import g_askdalle
from cmd_ai.g_askvision import g_askvision
from cmd_ai.syscom import process_syscom
from cmd_ai.version import __version__
from cmd_ai import best_examples

# from cmd_ai import best_examples

from cmd_ai.speak import str_to_mp3
import select
import click

# ===============================================================================================


def mark_code(text, colors ):
    """
    1.st it set cyan and default marks.
    2.nd it saves the code to /tmp
    """
    global count
    global code_buffer
    global PIPER

    #print(text)
    #print(text)

    # Define the pattern to match the pair of symbols
    pattern = "\s*```"

    # Define a counter to keep track of occurrences
    counter = 0

    # # Define a function to handle the replacement
    def replace(match ):
        nonlocal counter
        nonlocal colors  # from the mother function
        #print("...MATCH                   ",counter)
        counter+= 1
        if colors:
            bs=f"\n{fx.italic}{bg.default}{fg.lightcyan}```"    # code is in cyan
            es=f"\n```{fx.default}{bg.default}{fg.lightyellow}" # response is in yellow
        else:
            bs="\n#+begin_src "
            es="\n#+end_src"
        if counter % 2 == 1:
            return f"{bs}"   # return beginsource

            #return f"\n{fx.italic}{bg.default}{fg.lightcyan}{bs} "
            #return f"{fx.italic}{bg.default}{fg.lightcyan}#+begin_src python :results replace output :session test :exports results "
        else:
            return f"{es}"  # return endsource
            #return f"\n{es}{fx.default}{bg.default}{fg.lightyellow}"
            #return "{bg.default}{fg.default}"

    #.................................................................START
    # Use re.sub() with the defined function to replace the occurrences
    new_text = re.sub( pattern ,replace, text)

    if colors:
        new_text = f"{fg.lightyellow}{new_text}{fg.default}"
    else:
        new_text = f"{new_text}"
    return new_text


# ===============================================================================================


def display_response(res):
    """
    not only display in yellow, but extract eventual code and sace conversations.
    """
    resip = mark_code(res, colors = True)  # colors for terminal
    resiw = mark_code(res, colors = False) # good for emacs

    # I always clear the script
    config.PYSCRIPT_EXISTS = False
    config.SHSCRIPT_EXISTS = False

    # pythonista
    if config.CONFIG['current_role'] == "pythonista":
        if resiw.find("#+begin_src")>=0 and  resiw.find("#+end_src")>10:
            resco = re.sub(  r'^.*#\+begin_src(.*?)\n', "",  resiw, flags=re.DOTALL ) # ignore \n
            resco = re.sub( r'#\+end_src.*$',  "", resco , flags=re.DOTALL) # ignore \n
            with open( config.CONFIG['pyscript'] , "w" ) as f:
                f.write("#!/usr/bin/env python3\n\n")
                f.write(resco)
            config.PYSCRIPT_EXISTS = True

    # sheller
    if config.CONFIG['current_role'] == "sheller" or config.CONFIG['current_role'] == "piper":
        if resiw.find("#+begin_src")>=0 and  resiw.find("#+end_src")>10:
            resco = re.sub(  r'^.*#\+begin_src(.*?)\n', "",  resiw, flags=re.DOTALL ) # ignore \n
            resco = re.sub( r'#\+end_src.*$',  "", resco , flags=re.DOTALL) # ignore \n
            with open( config.CONFIG['shscript'] , "w" ) as f:
                f.write("#!/bin/bash\n\n")
                f.write(resco)
            config.SHSCRIPT_EXISTS = True


    # sourcecode ..........................................................
    #if config.CONFIG['current_role'] == "sheller":
    if resiw.find("#+begin_src")>=0 and  resiw.find("#+end_src")>10:

        match = re.search(r"#\+begin_src (\w+)", resiw, flags=re.DOTALL)
        if match:
            config.CONFIG['sourcecodeext'] = match.group(1)


        resco = re.sub(  r'^.*#\+begin_src(.*?)\n', "",  resiw, flags=re.DOTALL ) # ignore \n
        resco = re.sub( r'#\+end_src.*$',  "", resco , flags=re.DOTALL) # ignore \n
        OUTFILE = f"{config.CONFIG['sourcecode']}.{config.CONFIG['sourcecodeext']}"
        print( " ... written to ... ",OUTFILE )
        with open( OUTFILE , "w" ) as f:
            #f.write("#!/bin/bash\n\n")
            f.write(resco)
        config.SOURCECODE_EXISTS = True


    # piper can get python question...
    if config.CONFIG['current_role'] == "piper":
        if resiw.find("#+begin_src python")>=0 and  resiw.find("#+end_src")>10:
            resco = re.sub(  r'^.*#\+begin_src(.*?)\n', "",  resiw, flags=re.DOTALL ) # ignore \n
            resco = re.sub( r'#\+end_src.*$',  "", resco , flags=re.DOTALL) # ignore \n
            with open( config.CONFIG['pyscript'] , "w" ) as f:
                f.write("#!/usr/bin/env python3\n\n")
                f.write(resco)
            config.PYSCRIPT_EXISTS = True
        elif (resiw.find("#+begin_src bash")>=0 or resiw.find("#+begin_src sh")>=0) and  resiw.find("#+end_src")>10:
            resco = re.sub(  r'^.*#\+begin_src(.*?)\n', "",  resiw, flags=re.DOTALL ) # ignore \n
            resco = re.sub( r'#\+end_src.*$',  "", resco , flags=re.DOTALL) # ignore \n
            with open( config.CONFIG['pyscript'] , "w" ) as f:
                f.write("#!/bin/bash\n\n")
                f.write(resco)
            config.SHSCRIPT_EXISTS = True



    print( resip ) # fg.yellow, res, fg.default)
    #print( resiw ) # what will be in conversations.org
    # print(fg.yellow, res, fg.default)

    if config.CONFIG['current_role'] == 'pythonista' and config.PYSCRIPT_EXISTS:
        print(f"i...                        {fg.lightgreen}  see {config.CONFIG['pyscript']}; run with .e {fg.default}")

    elif config.CONFIG['current_role'] == 'sheller' and  config.SHSCRIPT_EXISTS:
        print(f"i...                        {fg.lightgreen}  see {config.CONFIG['shscript']}; run with .e {fg.default}")

    elif config.CONFIG['current_role'] == 'piper' and  config.SHSCRIPT_EXISTS:
        print(f"i...                        {fg.lightgreen}  see {config.CONFIG['shscript']}; run with .e {fg.default}")

    elif config.SOURCECODE_EXISTS:
        print(f"i...                        {fg.lightgreen}  see {config.CONFIG['sourcecode']}.{config.CONFIG['sourcecodeext']}; process with .e {fg.default}")


    # CREATE ORG FILE WITH THE RECORD
    if not os.path.exists(os.path.expanduser( config.CONFIG['conversations']) ):
        with open( os.path.expanduser( config.CONFIG['conversations']),"a") as f:
            f.write( texts.org_header )
            f.write("\n")

    # keep track with ORG
    with open( os.path.expanduser( config.CONFIG['conversations']),"a") as f:
        tnow = dt.datetime.strftime(dt.datetime.now(),'%Y/%m/%d %H:%M:%S')
        f.write(f"*** {config.CONFIG['last_prompt']}\n")
        # f.write(f" /{tnow} temperature={temp}/\n\n")
        f.write(f" /{tnow} /\n\n")
        f.write(resiw) # emacs oriented
        f.write("\n\n")

    # save last ( for mp3, speak)
    with open( os.path.expanduser( config.CONFIG['last_response']),"w" ) as f:
        f.write(f"{resiw}\n")
    ##########################################
    # READ# ./vety.py one  "`cat /tmp/cmd_ai_last_response.txt`" 3 -r
    modu = config.READALOUD%len(config.READALOUDSET)
    #print(f"i...  {bg.green}{fg.white} Reading Aloud is {config.READALOUDSET[modu]} {bg.default}{fg.default}")
    if config.READALOUDSET[modu] is not None:
        if config.SOURCECODE_EXISTS or config.PYSCRIPT_EXISTS or config.SHSCRIPT_EXISTS:
            print("x... CODE present .... no reading...sorry")
        else:
            str_to_mp3( resiw,  readme=True, lang=config.READALOUDSET[modu])



def pipe_mode():
    """
    detect mode - pipe or not
    """
    # this was blocking the detection of pipe...
    #if question is None or question=="" or len(question)<5: return None
    # LET US CONSIDER IT IS PIPE:
    #print(f"D... {fg.lightgray}Pipe mode OR NOT???? {fg.default}")
    if select.select([sys.stdin, ], [], [], 0.0)[0]:
        print(f"D... {fg.lightslategray}Pipe mode detected! {fg.default}")

        mstdin = []
        for line in sys.stdin:
            mstdin.append( line.strip() )
        return mstdin
    # UNCLESS IT IS NOT
    else:
        # print(f"D... {fg.lightgray} NO Pipe  {fg.default}")
        return None #print("No pipe")


###################################################################
#####################################################
###########################################
######################################

def main(cmd=None, budget=False, debug=False,
         u_switch = False, g_switch=False, csspeak=False, enspeak=False,
         version = False,
         multilineinput = False,
         limit_tokens = None,
         assistant = None,
         examples = False):
    """
chatGPT commandline interface

 it can run
 - interactive
 - as a standalone command ( cmd_ai "Say Hi")
 - in pipe ( echo "Say Hi" | cmd_ai -a  pythonista) ; default assistant is Ubuntu 22 expert
Parameters
  assistant :   ... assistant name: assistant or None | translator | sheller | pythonista | dalle
  limit_tokens : ...  limit  tokens. 300 is very basic -l
  multilineinput : ... tricky, use Alt-Enter instead of enter. Do not use in pipe
  budget : ... print the budget on this pc, split by month and by day in the last month
  u_switch : ... add ability - spectial utils - testing
  g_switch : ... add google search 1st five ability
  version : ... program version
  examples : ... show examples
"""
    #                        HELP help
    if examples:
        best_examples.main()
        return 0
    if version:
        print(__version__)
        return 0
    # ======== DEFINE THE CONFIG FILE HERE ========


    if os.path.exists(  os.path.expanduser(config.CONFIG['pricelog'] ) ):
        # dateparse = lambda x: dt.datetime.strptime(x, '%Y-%m-%d %H:%M:%S')
        df = pd.read_csv( os.path.expanduser(config.CONFIG['pricelog']) ,
                          delim_whitespace=True,
                          names = ['date','time', 'in','out','price']
                         )#, parse_dates=['datetime'], date_parser=dateparse)

        df['datetime'] = pd.to_datetime(df.date)
        df['year'] = pd.to_datetime(df.date).apply(lambda x:x.year)
        df['month'] = pd.to_datetime(df.date).apply(lambda x:x.month)
        #df.index = df['datetime']
        #df['priced'] = df.groupby(df['datetime'])['price'].sum()
        df1 = df.groupby('month', sort=False).agg({'datetime':'last','price':'sum'})# , 'in':'sum','out':'sum'
        # first by mothn
        df = df.groupby('datetime',sort=False).agg({'price':'sum', 'in':'sum','out':'sum','year':'first','month':'first'})

        if budget:
            print("____________________________________")
            print(df1)#.iloc[-30:-1,:]  )
            print("____________________________________")
        print(df.loc[ df['month']==2 ]  )#.iloc[-30:-1,:]  )
        print("____________________________________")

    if budget: return

    config.CONFIG["filename"] = "~/.config/cmd_ai/cfg.json"
    config.load_config()

    # ================================= changes to config FROM HERE ==========
    if limit_tokens is not None:
        config.CONFIG['limit_tokens'] = limit_tokens
    print("D... init tokens",config.CONFIG['limit_tokens'] ,  limit_tokens )


    if config.CONFIG["api_key"] is None:
        res = get_api_key()
        if res is not None:
            config.CONFIG["api_key"] = res
            res = input("> save api_key to config?   y/n")
            if res == "y":
                print(config.CONFIG["api_key"])
                config.save_config()
            else:
                print(" ... not now ...")


    config.tokens = 0
    config.started_task = dt.datetime.now()
    config.started_total = dt.datetime.now()

    # in CON config.PYSCRIPT = "/tmp/gpt_code.py"
    config.PYSCRIPT_EXISTS = False
    config.SHSCRIPT_EXISTS = False


    config.pipeinput = pipe_mode()
    if config.pipeinput is not None:
        pass
        # system prompt must be linux expert

    config.client = OpenAI(api_key=get_api_key())
    config.myPromptSession = PromptSession(
        history=FileHistory(os.path.expanduser("~/.cmd_ai.history")) #, multiline=True
    )

    # config.load_config()
    # config.save_config()

    if cmd == "usage":
        print(
            """ ... usage:
                _
        """
        )
        sys.exit(0)
    if cmd == "savecfg":
        config.save_config()
        print("i... config saved")
        sys.exit(0)
    # else:
    #    unitname.func()

    # IN SYSCOM
    # models = config.client.models.list()
    # mids = []
    # for i in models.data:
    #     if i.id.find("gpt") >= 0:
    #         mids.append(i.id)
    # for i in sorted(mids):
    #     print("   ", i)

    print("                            ", fg.cyan, ".h for help", fg.default)

    # DECIDE WHETHER PIPE .....

    # if config.pipeinput is not None:
    #     print("=========== PIPE ========")
    # else:
    #     print("=========== RUN INTERACTIVE ========")
    # sys.exit(0)

    if config.pipeinput is not None and multilineinput:
        print(fg.red, "X... no multiline input allowed for PIPE mode") # it crashes
        sys.exit(1)

    if (config.pipeinput is None) and (cmd is None): #if cmd == "":
        # print("========================= NONE PIPE MODE ===================")
        while True:

            SYSCOM = False
            inp = config.myPromptSession.prompt("> ", multiline=multilineinput)
            config.started_task = dt.datetime.now()

            #print("D:1")
            inp = inp.strip()
            if config.CONFIG["current_role"] == "pythonista" and  inp[0]!="." and len(inp)>2:
                inp = f"{inp} (Write python code)."
            if config.CONFIG["current_role"] == "sheller" and  inp[0]!="." and len(inp)>2:
                inp = f"{inp} (Write bash code if needed)."


            if inp.strip() == "":
                continue
            if len(inp.strip()) == 1:
                print(f"!... {fg.red}  one letter prompt not allowed, use .h  {fg.default}")
                continue

            if len(inp) > 0 and len(inp) < 10 and inp[0] == ".":
                # System command .... root style ...
                SYSCOM = True
                print(fg.orange, "command: ",inp)
            else:
                config.CONFIG['last_prompt'] = inp
                #print(fg.white, inp)
            print(fg.default)

            if SYSCOM:
                process_syscom(inp)
            elif config.CONFIG["current_role"] == "dalle":
                print("D... in DALLE role ...")
                res = g_askdalle(inp)
                print("D...  done:", res)
            else:
                # ================================================vision OR gpt
                stoplength = False
                if config.CONFIG["current_role"] == "vision":
                    print("D... in VISION role ...")
                    res, stoplength = g_askvision(inp)
                    #print("D...  done:", res)
                else:
                    #print("D:2")
                    res,stoplength = g_askme(inp, model="gpt-4-1106-preview")
                # --------------------- and I continue as normal--------------
                # append the response...
                config.messages.append( {"role":"assistant", "content":res} ) # always user/assistant
                # this I want to be json
                with open( os.path.expanduser( config.CONFIG['current_messages'] ), "w" )  as f:
                    try:
                        f.write( json.dumps( config.messages, indent=2, separators=(',', ': ')) )
                    except:
                        print(fg.red,f"X... I cannot write {len(config.messages)} messages to json {config.CONFIG['current_messages']}. web content? ")
                    #f.write( str(config.messages) )

                display_response(res)
                if stoplength:
                    print(f"!... {fg.red} stopped because : not enough tokens, try .l {fg.default}")
    else:# cmd is not None:
        # print("========================= fire and forget mode ===================")
        # commandline.  OR pipe on commandline ====================================
        #  this gets more complex with --assistant  SWITCH
        #

        inp = cmd # command
        if inp is None: inp = "" # protection for pipe and different assistant
        # for cmdline, I operate the switches here
        if u_switch:
            process_syscom(".u")
        elif g_switch:
            process_syscom(".g")

        if csspeak:#_switch:
            config.READALOUD = 2 #
        elif enspeak:#_switch:
            config.READALOUD = 1 #
        #print("cs=",vcs,"en=",ven)
        #str_to_mp3( "Žluťoučký kůň",  readme=True, lang=config.READALOUDSET[config.READALOUD%len(config.READALOUDSET)])

        if assistant is None:
            print(f"i...  {bg.green}{fg.white} Ubuntu expert {bg.default}{fg.default}")
            config.messages.append({"role": "system", "content": texts.role_piper})
            config.CONFIG["current_role"] = "piper"
        else:
            print(f"i...  {bg.green}{fg.white} {assistant}  {bg.default}{fg.default}")
            if assistant == "translator":
                config.messages.append({"role": "system", "content": texts.role_translator})
            elif assistant == "assistant":
                config.messages.append({"role": "system", "content": texts.role_assistant})
            elif assistant == "sheller":
                config.messages.append({"role": "system", "content": texts.role_sheller})
            elif assistant == "pythonista":
                config.messages.append({"role": "system", "content": texts.role_pythonista})
            elif assistant == "dalle":
                config.messages.append({"role": "system", "content": texts.role_dalle})
            elif assistant == "vision":
                config.messages.append({"role": "system", "content": texts.role_vision})
            else:
                print("X... NOT KNOWN ASSISTANT:  translator | assistant | sheller | pythonista | piper | dalle | vision")
                sys.exit(1)
            config.CONFIG["current_role"] = assistant


        if config.pipeinput is not None: # pipe + cmd
            print(f"i...  {bg.green}{fg.white} Using PIPE {bg.default}{fg.default}")
            pipe_input = '\n'.join(config.pipeinput)
            # print(pipe_input)
            if assistant is None:
                # (Write bash code if required). ....  ??
                inp = f"Here is some linux command output:\n```{pipe_input}```\n{inp}"
            else:
                # I hope that assistant knows what to do
                if inp is not None or len(inp)> 0:
                    print(fg.red,f"i... disregarding text: /{inp}/", fg.default)
                inp = pipe_input
        #print(fx.italic,inp, fx.default,"\n")

        if assistant != "dalle":

            res = ""
            stoplength = False
            # --------------------------------------------------------
            if assistant == "vision":
                res,stoplength = g_askvision(inp)
            else:
                res,stoplength = g_askme(inp, model="gpt-4-1106-preview")

            #= ----------------------------------------------
            config.messages.append( {"role":"assistant", "content":res} ) # always user/assistant

            with open( os.path.expanduser( config.CONFIG['current_messages'] ), "w" )  as f:
                f.write( str(cmd) )
            display_response(res)
            if stoplength:
                print(f"!... {fg.red} stopped because : not enough tokens, try .l {fg.default}")

            if config.SHSCRIPT_EXISTS:
                click.echo('Run the code? [yn] ', nl=False)
                c = click.getchar()
                click.echo()
                if c == 'y':
                    click.echo('We will go on')
                    process_syscom(".e")
                elif c == 'n':
                    click.echo('Abort!')
                else:
                    click.echo('not n nor y : Abort ')
                #if click.confirm('Run the code?', abort=True, default=False):
                #    process_syscom(".e")

        else:
            res = g_askdalle(inp)



# ====================================================================

if __name__ == "__main__":
    Fire(main)
