#!/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 calc_budget
from cmd_ai import operate_context_files
# from cmd_ai import best_examples

from cmd_ai.speak import str_to_mp3
import select
import click

# ===============================================================================================
MODEL_TO_USE = "gpt-4-turbo-2024-04-09"
MODEL_TO_USE = "gpt-4o-2024-05-13"

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 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
        operate_context_files.save_source_code( resco )
        # 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( operate_context_files.get_filename_my_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( operate_context_files.get_filename_my_shscript() ) 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}  see1 {operate_context_files.get_filename_my_pyscript()}; run with .e {fg.default}")

    if config.CONFIG['current_role'] == 'sheller' and  config.SHSCRIPT_EXISTS:
        print(f"i...                        {fg.lightgreen}  see2 {operate_context_files.get_filename_my_shscript()}; run with .e {fg.default}")

    if config.CONFIG['current_role'] == 'piper' and  config.SHSCRIPT_EXISTS:
        print(f"i...                        {fg.lightgreen}  see3 {operate_context_files.get_filename_my_shscript()}; run with .e {fg.default}")

    if config.SOURCECODE_EXISTS:
        print(f"i...                        {fg.lightgreen}  see4 {operate_context_files.get_filename_my_sourcecode()}; 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)
    operate_context_files.save_last_response(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.lightslategray}Pipe mode NOT detected .... producing more output {fg.default}")
        # print(f"D... {fg.lightgray} NO Pipe  {fg.default}")
        return None #print("No pipe")


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

def main(cmd=None,
         assistant = None, # role, predefined assistant system prompt
         budget=False, #  CLI ... show budget
         debug=False,
         csspeak=False, #
         enspeak=False, #
         examples = False, # show
         g_switch=False, # google search function
         u_switch = False, # weather function
         limit_tokens = None, # default is very small 300
         multilineinput = False, # enter will be newline, altenter? is enter
         name = None, # config['zulu'] default name of the AI
         reset = False, #  for commandline PIPE - forget totaly the context
         version = False,
         zzz = None
         ):
    """
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 ========

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

    if name is None:  # zulu is default in config
        name = config.CONFIG['current_name']
    else:
        config.CONFIG['current_name'] = name

    # ================================= changes to config FROM HERE ==========
    if limit_tokens is not None:
        config.CONFIG['limit_tokens'] = limit_tokens

    # First - tell if pipe mode: ------------------------
    config.pipeinput = pipe_mode()  # detect pipe mode         # system prompt must be linux expert
    if config.pipeinput is None:

        # PRICING AND BUDGET
        calc_budget.show_budget( budget )
        if budget: return

        print("D... initial  token limit:",config.CONFIG['limit_tokens']   )
        print("                            ", bg.cyan, ".h for help;  .r for reset", bg.default)


    # ======================================= recover api token and save to 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 ...")

    # =============== global variables
    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.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 == "savecfg":
        config.save_config()
        print("i... config saved")
        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)


    # and if reset => destroy all
    if reset:
        print(f"i... I am reseting all history...  from the next question, alle is lost")
        config.messages = []  # i hope it automaticall assigns a role
        config.PYSCRIPT_EXISTS = False
        config.SHSCRIPT_EXISTS = False
        totsize,totlines = 0, 0
    else:
        totsize,totlines = operate_context_files.load_all_config_messages()


    # ***********************************************************************************  MAIN INTERACTIVE
    # ***********************************************************************************  MAIN INTERACTIVE
    # ***********************************************************************************  MAIN INTERACTIVE
    if (config.pipeinput is None) and (cmd is None):
        print(f"i... My name is {fg.orange}{config.CONFIG['current_name']}{fg.default} and I remember {fg.orange}{totlines}{fg.default} lines from previous conversations.")

        # 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:
                    res,stoplength = g_askme(inp, model=MODEL_TO_USE)    # *************************************

                # --------------------- and I continue as normal--------------append response and SAVE
                config.messages.append( {"role":"assistant", "content":res} ) # always the pair:  user/assistant
                operate_context_files.save_all_config_messages()
                # # append the response...
                # # 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:  -----------------------------------------------------------  PIPE ----------------
        # cmd is not None:  -----------------------------------------------------------  PIPE ----------------
        # cmd is not None:  -----------------------------------------------------------  PIPE ----------------
        # cmd is not None:  -----------------------------------------------------------  PIPE ----------------
        # cmd is not None:  -----------------------------------------------------------  PIPE ----------------
        # print("========================= fire and forget mode ===================")
        # commandline.  OR pipe on commandline ====================================
        #  this gets more complex with --assistant  SWITCH
        #

        inp = cmd # inp is the text comming with the cmd ...... echo 'pipe_input' | cmd_ai 'inp'
        pipe_input = '\n'.join(config.pipeinput)
        config.silent = True # omit lot of messages

        if inp is None: inp = "" # protection for pipe and different assistant

        if u_switch: # for cmdline, I operate the switches here
            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: # means default piper
            #print(f"i...  {bg.green}{fg.white} Ubuntu expert {bg.default}{fg.default}")
            if len(config.messages)>0:
                config.messages[0] = {"role": "system", "content": texts.role_piper}
            else:
                config.messages.append({"role": "system", "content": texts.role_piper})
            config.CONFIG["current_role"] = "piper"
            inp = f"Here is a linux command output:\n```{pipe_input}```\n{inp}"
        else:
            print(f"i...  {bg.green}{fg.white} assistant role override: {assistant}  {bg.default}{fg.default}")
            if assistant == "translator":
                if len(config.messages)>0:
                    config.messages[0]={"role": "system", "content": texts.role_translator}
                else:
                    config.messages.append({"role": "system", "content": texts.role_translator})
            elif assistant == "assistant":
                if len(config.messages)>0:
                    config.messages[0]={"role": "system", "content": texts.role_assistant}
                else:
                    config.messages.append({"role": "system", "content": texts.role_assistant})
            elif assistant == "sheller":
                if len(config.messages)>0:
                    config.messages[0]={"role": "system", "content": texts.role_sheller}
                else:
                    config.messages.append({"role": "system", "content": texts.role_sheller})
            elif assistant == "pythonista":
                if len(config.messages)>0:
                    config.messages[0]={"role": "system", "content": texts.role_pythonista}
                else:
                    config.messages.append({"role": "system", "content": texts.role_pythonista})
            elif assistant == "dalle":
                if len(config.messages)>0:
                    config.messages[0]={"role": "system", "content": texts.role_dalle}
                else:
                    config.messages.append({"role": "system", "content": texts.role_dalle})
            elif assistant == "vision":
                if len(config.messages)>0:
                    config.messages[0]={"role": "system", "content": texts.role_vision}
                else:
                    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
            inp = f"Here is an input:\n```{pipe_input}```\n{inp}"


        # if config.pipeinput is not None: # pipe + cmd
        #     print(f"i...  {bg.green}{fg.white} Using PIPE {bg.default}{fg.default}")
        #     # 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=MODEL_TO_USE)

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

            curmessages = operate_context_files.get_filename_my_context()
            with open( os.path.expanduser( curmessages ), "w" )  as f:
                f.write( json.dumps( config.messages, indent=2, separators=(',', ': ')) )
                # f.write( config.messages ) #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)
