"""
Implementation of the completion-based capabilities in openai_gpt.py, used for
AI suggest and transform.

Copyright (c) 1999-2024, Archaeopteryx Software, Inc.  All rights reserved.

Written by Stephan R.A. Deibel and John P. Ehresman
"""

import os
import sys

# Replaced by code at the end of openai_gpt.py:
def _OverrideConfig(config):
    raise NotImplementedError('should be replaced')
def _APIKey():
    raise NotImplementedError('should be replaced')
def _GetConfig(key):
    raise NotImplementedError('should be replaced')
_SUPPORTED_PROVIDER_BASE_URLS = None

SUGGEST_PROMPT = '''
This automated call is being made by a plugin to Wing IDE that allows you to help the user with their code. Therefore your reply must not contain any preambles or other explanations, because it'll make the plugin code break.
{files}
The user currently has the file `{filename}` open and is requesting suggestions for code to place at the current insertion point (`insertion_position`) located at character {insertion_position} in the editor.

This is the code that precedes the insertion point in the editor (`preceding_code`):

```
{preceding_code}
```

Below is the code the follows the insertion point in the editor (`following_code`):

```
{following_code}
```
{request}
Your reply (`your_reply`) should be a single code block correctly indented for the target position, with no explanations, delimited by triple backticks. We'll perform `new_code = preceding_code + your_reply + following_code`, which means that your code will be added to the file at the current insertion point.
'''

SUGGEST_FILES = '''
The following additional files may be relevant to providing suggestions to the user.  These are files that are being imported into the currently opened file or are in the user's project:
'''

SUGGEST_REQUEST = '''
The user has provided the following message / request for the code they would like you to suggest for the current insertion point (`insertion_position`):

{request}
'''

TRANSFORM_PROMPT = '''
This automated call is being made by a plugin to Wing IDE that allows you to help the user with their code. Therefore your reply must not contain any preambles or other explanations, because it'll make the plugin code break.
{files}
The user currently has the file `{filename}` open and is requesting you to change the selected code (`selected_code`) located at character {selection_position} in the editor.

This is the user's message / request (`user_message`):

```
{request}
```

Below is the code in the currently open file (`entire_file_content`):

```
{text}
```

Below is the code that's selected (`selected_code`). The user's message / request pertains only to this segment of code, and it's the only segment you're allowed to change:

```
{selected_text}
```

Your reply should be a single code block with no explanations, delimited by triple backticks.  Please include all unaltered lines in your reply, including also blank lines in parts of the code you have not changed.  We'll perform `new_code = entire_file_content.replace(selected_code, your_reply)`, which means that the selected code will be deleted and replaced with your code.
'''

TRANSFORM_PROMPT_NEW_FILE = '''
This automated call is being made by a plugin to Wing IDE that allows you to help the user with their code. They are asking you to generate a response that will in its entirety be pasted into a new editor.
{files}
The user currently has the file `{filename}` open and is focusing on the selected code (`selected_code`) located at character {selection_position} in the editor.

This is the user's message / request (`user_message`):

```
{request}
```

Below is the code in the currently open file (`entire_file_content`):

```
{text}
```

Below is the code that's selected (`selected_code`). The user's message / request pertains only to this segment of code:

```
{selected_text}
```

Your reply should be a single block of text delimited by triple backticks.  Please indicate, immediately after the opening triple backticks, the file type that is suggested by the user's request.
'''

INFO_PROMPT = '''
This automated call is being made by a plugin to Wing IDE that allows you to help the user with their code. They are asking you to assist them with a coding task.
{files}
The user is focusing on and asking about the file `{filename}`.

This is the user's message / request (`user_message`):

```
{request}
```

Below is the contents of the file they are focusing on (`entire_file_content`):

```
{text}
```

Your reply should be a single block of text delimited by triple backticks.  Please indicate, immediately after the opening triple backticks, the file type that for the text of your response.
'''

TRANSFORM_FILES = '''
The user has identified the following additional files as relevant to their request:
'''

CHAT_VIA_COMPLETIONS_API_PROMPT = '''
The user currently has the file `{filename}` open and is focusing on the selected code (`selected_code`) in the editor.

Below is the content of the currently open file:

```
{text}
```

Below is the code that's selected (`selected_code`):

```
{selected_text}
```
'''

CHAT_VIA_COMPLETIONS_API_PROMPT_NO_SELECTION = '''
The user currently has the file `{filename}` open with the following content:

```
{text}
```
'''

def _RunQueryProcess(messages: list, result_cb) -> str:
    """Run a query using the configured API key and language model. The (exit_code, error_output,
    cost, result) is passed to the given result_cb asynchronously. This starts a sub-process to
    satisfy the request."""

    import _winghome
    from wingbase import spawn
    import bootstrap_utils
    
    pyexe = _GetPyExec()
   
    args = [
        pyexe,
        os.path.abspath(__file__),
    ]
    
    env = os.environ.copy()
    env['WINGHOME'] = _winghome.kWingHome
    env['WINGSETTINGSDIR'] = _winghome.kUserSettings
    
    # The plugin can self-install openai but this is no longer needed because Wing
    # ships with openai now.  However, an env can be set to enable that and then
    # want to run using the self-installed openai.  See also notes in ai/manager.pydantic
    # in CAIManager.EnsureAIProviderDepdencies:  That and the installation support
    # here is slated to be moved into a facility available to all plugin authors.
    if 'WING_ENABLE_OPENAI_INSTALLER' in os.environ:
        from _openai_installer import _FindInstallDirs
        env['PYTHONPATH'] = _FindInstallDirs()[1]
        
    # Need to add the external-packages directory to the path so openai can be imported
    else:
        platform = bootstrap_utils.NP_GetPlatform()
        ext_pkgs = os.path.join(_winghome.kWingHome, 'bin', '__os__', platform,
                                'external-packages')
        env['PYTHONPATH'] = ext_pkgs
        
    env['PYTHONOPTIMIZE'] = '2'
    env['PYTHONIOENCODING'] = 'utf-8'

    print("AI Query", args)
    proc = spawn.CChildProcess(args, env=env, child_pwd=os.getcwd(), separate_stderr=True,
                               io_encoding='utf-8')
    
    def start_failed(obj, error):
        
        result_cb(-1, error, 0, None)
        
    def terminated(obj):
        
        exit_code = proc.GetExitCode()
        error = proc.GetStdErrOutput()
        result = proc.GetOutput()

        if result:
            lines = result.splitlines(keepends=True)
            assert len(lines) >= 3
            cost = (
                _GetConfig('currency'),
                int(lines[0].strip()) / 1000 * _GetConfig('prompt_tok') + int(lines[1].strip()) / 1000 * _GetConfig('result_tok')
            )
            result = ''.join(lines[2:])
        else:
            cost = (_GetConfig('currency'), 0., 0.)
            result = ''
            
        result_cb(exit_code, error, cost, result)
        
    proc.connect('terminated', terminated)
    proc.connect('start-failed', start_failed)
    
    proc.Start()
    
    brand, base_url = _GetConfig('base_url')
    if brand not in ('custom', 'azure'):
        base_url = _SUPPORTED_PROVIDER_BASE_URLS[brand]
    if base_url is None:
        base_url = 'None'
    mlines = []
    for mtype, qtype, msg in messages:
        mlines.append('|||')
        mlines.append(mtype)
        mlines.extend(msg.splitlines())
    inp = brand + '\n' + base_url + '\n' + _APIKey() + '\n' + _GetConfig('model') + '\n' + '\n'.join(mlines)
    proc._SendInputToChild(inp, close=True)
	
    description = [
        spawn._win32_cmd_line_from_argv(args),
        'WINGHOME={}'.format(_winghome.kWingHome),
        'WINGSETTINGSDIR={}'.format(_winghome.kUserSettings),
        'Query:',
    ]

    for mtype, qtype, msg in messages:
        description.append(mtype.title()+': '+msg)
        
    description = '\n'.join(description)
    
    return proc, description

def _GetPyExec():

    import bootstrap_utils
    import _winghome
    
    # Don't use sys.executable because it is not Python when Wing is running in production
    pyver = '{}.{}'.format(sys.version_info.major, sys.version_info.minor)
    pydir = os.path.join(_winghome.kWingHome, 'bin', '__os__', bootstrap_utils.NP_GetPlatform(),
                             'runtime-python{}'.format(pyver))
    
    if sys.platform == 'win32':
        pyexe = os.path.join(pydir, 'python.exe')
    else:
        pyexe = os.path.join(pydir, 'bin', 'python{}'.format(pyver))

    return pyexe

def CompletionQuery(query: str, result_cb) -> str:
    """Run the given query string through OpenAI compatible completion API. This calls the given
    result_cb with the (error, cost, response) where error is a non-empty error on failure
    and otherwise response is the code suggested by the language model. Return value is
    the spawn.CProcess object running the query and a description of the query for
    diagnostic purposes."""
    
    import _openai_chat
    
    def query_complete(exit_code, error, cost, response):

        if exit_code != 0:
            error = 'Process failed with exit_code={}; output:'.format(exit_code) + '\n' + error

        mime_type = None
        
        if response is not None:
            try:
                
                result_start = response.find('```')
                if result_start != -1:
                    
                    result_start += 3
                    ticks_end = result_start
                    found_eol = False
                    while not found_eol and result_start < len(response):
                        if response[result_start:result_start+2] == '\r\n':
                            result_start += 2
                            found_eol = True
                        elif response[result_start] in '\r\n':
                            result_start += 1
                            found_eol = True
                        else:
                            result_start += 1
                    file_type = response[ticks_end:result_start].strip()
                    mime_type = _openai_chat._kIndicatorMimeTypes.get(file_type, 'text/plain')
                    
                    result_end = response.rfind('```')
                    if result_end == result_start:
                        result_end = len(response)
                        
                    if response[result_end-2:result_end] == '\r\n':
                        result_end -= 2
                    elif response[result_end-1] in '\r\n':
                        result_end -= 1
                    
                    response = response[result_start:result_end]

            except:
                errstr = "Invalid response from AI Provider: '{}'".format(response)
                if error:
                    errstr += '\n\n' + "Error:" + '\n' + error
                error = errstr

        if error:
            print("AI provider error output with apparently valid response:")
            print(error)
            print(response)
            
        result_cb(error, cost, response, mime_type)
        
    messages = [('user', 'text', query)]
    
    return _RunQueryProcess(messages, query_complete)

def _EstimateTokens(query: str) -> int:
    """Count the approximate number of tokens a query will consume, in order to keep
    usage below configured thresholds.  This is not used in production, but useful
    during development."""
    
    # This is an approximation
    return len(str) / 5

    # Getting an exact value would require us to include tiktoken with Wing, including
    # also all its many dependencies, and is probably not worth doing that. It downloads
    # the encoding the first time it's run, which further complicates our licensing and
    # security picture.
    #import tiktoken
    #encoding = tiktoken.encoding_for_model(_GetConfig('model')
    #return len(encoding.encode(query))
    
def GetFilesPart(files):
    
    if not files:
        return ''
    
    filepart = []
    for loc in files:
        if files[loc] is None:
            continue            
        content, mime_type, description = files[loc]
        if content is None:
            continue
        content = content.replace('```', '"""')
        
        description = description.replace("'", '"').replace('```', '"""')
        
        filepart.append('')
        if loc.GetUrl().startswith('unknown:'):
            filepart.append(f"Text fragment with mime type {mime_type} described as '{description}':")
        else:
            filepart.append(f"File {loc._GetNativeFullPath()} with mime type {mime_type} described as '{description}':")
        filepart.append('```')
        filepart.append(content)
        filepart.append('```')
        filepart.append('')
        
    files = '\n'.join(filepart)

    return files
        
        
def CompletionsAPIMain():
    """This implements the 'server' used for calling the OpenAI compatible completions API for AI
    code complete and AI transform."""
    
    lines = sys.stdin.readlines()
    brand = lines[0].strip()
    base_url = lines[1].strip()
    if base_url == 'None':
        base_url = None
    key = lines[2].strip()
    model = lines[3].strip()

    messages = []
    new_item = False
    mtype = None
    cur_message = []
    for line in lines[4:]:
        if new_item:
            mtype = line.strip()
            new_item = False
        elif line.strip() == '|||':
            if cur_message:
                messages.append({'role': mtype, 'content': '\n'.join(cur_message)})
            mtype = None
            new_item = True
            cur_message = []
        elif mtype is not None:
            cur_message.append(line)
        else:
            raise RuntimeError('malformed input in completions server: '+line)

    if cur_message:
        messages.append({'role': mtype, 'content': '\n'.join(cur_message)})
        
    if not messages:
        raise RuntimeError('no messages found in completions server'+'\n'+''.join(lines[4:]))
    
    if brand == 'azure':
        result = _RunAzureOpenAICompletionRequest(key, base_url, model, messages)
    else:
        result = _RunOpenAICompletionRequest(key, base_url, model, messages)

    print(result.usage.prompt_tokens)
    print(result.usage.completion_tokens)
    print(result.choices[0].message.content.replace('\r\n', '\n'))
    
def _RunAzureOpenAICompletionRequest(key, base_url, model, messages):
    
    from openai import AzureOpenAI
    
    os.environ['AZURE_OPENAI_API_KEY'] = key
    
    client = AzureOpenAI(
        api_version=os.environ.get('OPENAI_API_VERSION', '2024-12-01-preview'), 
        azure_endpoint=base_url,
    )
    
    result = client.chat.completions.create(
        model=model, 
        messages=messages
    )
    
    return result
    
def _RunOpenAICompletionRequest(key, base_url, model, messages):
    
    from openai import OpenAI, resources
    
    client = OpenAI(
        api_key=key,
        base_url=base_url       
    )
    
    chat = resources.Chat(client)
    result = chat.completions.create(
        model=model,
        messages=messages,        
    )

    return result
    

#########################################################################
if __name__ == '__main__':
    CompletionsAPIMain()

