#!/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.syscom import process_syscom
from cmd_ai.version import __version__

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(question):
    """
    detect mode - pipe or not
    """
    if question is None or question=="" or len(question)<5: return None
    if select.select([sys.stdin, ], [], [], 0.0)[0]:
        #print(f"D... {fg.lightgray}Pipe mode! {fg.default}")

        mstdin = []
        for line in sys.stdin:
            mstdin.append( line.strip() )
        return mstdin
    else:
        return None #print("No pipe")


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

def main(cmd="", budget=False, debug=False, u_switch = False, g_switch=False, csspeak=False, enspeak=False, version = False):
    """
    chatGPT commandline portal: \n
    I can run interactive, standalone or in pipe
    --version program version
    --budget see the budget
    cmd ... Text fo chat in parentheses
    --u_switch add u ability
    --g_switch add google ability

"""
    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'})

        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()

    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(cmd)
    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"))
    )

    # 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)
    if cmd == "":
        while True:

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

            #print("D:1")
            inp = inp.strip()
            if config.CONFIG["current_role"] == "pythonista" and  inp[0]!=".":
                inp = f"{inp} (Write python code)."
            if config.CONFIG["current_role"] == "sheller" and  inp[0]!=".":
                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)
            else:
                #print("D:2")
                res,stoplength = g_askme(inp, model="gpt-4-1106-preview")

                # 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:
        # commandline.  OR pipe on commandline
        inp = cmd # command

        # 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)])

        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"

        if config.pipeinput is not None: # pipe + cmd
            nl = '\n'.join(config.pipeinput)

            # (Write bash code if required).
            inp = f"Here is some linux command output:\n```{nl}```\n{cmd}"
            print(f"i...  {bg.green}{fg.white} Pipe expert {bg.default}{fg.default}")

        #print(fx.italic,inp, fx.default,"\n")

        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('Invalid input :(')
            #if click.confirm('Run the code?', abort=True, default=False):
            #    process_syscom(".e")

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

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