"""
Implementation of OpenAI Responses API support, used for AI Chat.

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

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

import sys
import os
import threading
import traceback
import collections
import time
import hashlib
import io
import select
import json
import socket
import struct

MAX_ASSISTANT_CONTEXT_FILES = 20

def _LoadOpenAIGPTModule():
  old_path = list(sys.path)
  try:
    sys.path.append(os.path.dirname(os.path.abspath(__file__)))
    import openai_gpt
  finally:
    sys.path = old_path
    
  return openai_gpt

# The following are replaced by code at the end of openai_gpt.py when this module is imported by the
# plugin managed in the IDE to avoid importing openai_gpt in that case, because that results in a
# duplicate copy of the loaded module because the plugin manager does not add to sys.modules.
def _OverrideConfig(config):
  openai_gpt = _LoadOpenAIGPTModule()
  openai_gpt.CONFIG = config
def _APIKey():
  openai_gpt = _LoadOpenAIGPTModule()
  return openai_gpt._APIKey()
def _GetConfig(key):
  openai_gpt = _LoadOpenAIGPTModule()
  return openai_gpt._GetConfig(key)
  
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

gClient = [None]
def _GetClient():
  """Get the overall openai API client"""

  from openai import OpenAI

  if gClient[0] is None:
    gClient[0] = OpenAI(api_key=_APIKey())

  return gClient[0]


class RateLimit:
  def __init__(self, per_second=10):
    self.per_second = per_second
    self.last_call = time.time()
    
  def tick(self):
    sleep_time = (1.0 / self.per_second) - (time.time() - self.last_call)
    if sleep_time > 0.0:
      time.sleep(sleep_time)
    self.last_call = time.time()


class Assistant():
  """AI assistant class for AI Chat.  This implementation uses the OpenAI
  Responses API with file search via vector stores."""

  all_assistants = {}

  def __init__(self):
    self.__clear()

  def create(self, name):
    """Initialize the assistant, using the given assistant name.  This creates
    an OpenAI vector store for file search context.  Returns (id, error) where
    id is the vector store id on success and error is an error string on failure."""

    client = _GetClient()

    # Create vector store for file search context
    try:
      self.vector_store = client.vector_stores.create(
        name=('wing:'+name)[:255],
      )
    except:
      error = []
      error.append("Could not create vector store for AI Assistant:")
      error.append(traceback.format_exc())
      error = '\n'.join(error)
      return (None, error)

    self.model = _GetConfig('model')
    self.instructions = ("You are a programmer's assistant that helps a software developer "
                         "write, understand, debug, and modify code written in Python and other "
                         "programming and markup languages.")

    self.all_assistants[self.vector_store.id] = self

    return (self.vector_store.id, None)

  def delete(self):
    """Delete the vector store and all associated uploaded files from the
    user's OpenAI account. Returns a list of errors if there were any failures or None on
    success. If there are errors, some resources could not be deleted and should be
    removed later with static methods delete_assistant_from_account() and
    delete_uploaded_file_from_account()."""

    client = _GetClient()

    # Delete the vector store
    try:
      client.vector_stores.delete(self.vector_store.id)
    except:
      pass

    # Delete all threads
    for thread in list(self.threads.values()):
      thread.delete()
    self.threads = {}

    # Delete all uploaded files associated with this assistant from the user account
    errors = []
    for filename, content_hash in list(self.uploaded_files.keys()):
      file_id = self.uploaded_files[(filename, content_hash)][0]
      error = self.delete_uploaded_file_from_account(file_id)
      if error:
        errors.append(error)
    self.uploaded_files = {}

    if self.vector_store is not None:
      self.all_assistants.pop(self.vector_store.id, None)

    # Clear to blank even if resources were left, since removing them will
    # be retried later without this Assistant instance
    self.__clear()

    return errors

  def get_id(self):
    """Get the ID for this assistant (the vector store ID)"""

    return self.vector_store.id

  def set_context(self, files):
    """Set the context files to use for this assistant as a whole.  This includes
    all the files other than the current file that are relevant to conversation
    threads hosted by the assistant.  The files should be sent most important first
    since OpenAI limits how many can be associated with the assistant.
    
    The arg files is a list of (filename, content_file, description) where content_file
    is None if the primary file content should be used or non-None if it contains the
    content to use for the file, encoded as utf-8.

    Returns (file_ids, error) where file_ids is a list of Open AI file ids for the
    context files on success and error is and error message if one or more file
    failed to upload or associate with the assistant.

    When there are errors, there may still be items in file_ids for the files that were
    successfully processed."""

    errors = []

    # Maximum number of files imposed by OpenAI API
    files = files[:MAX_ASSISTANT_CONTEXT_FILES]

    old_context_file_ids = set(self.context_files)

    # Upload each file if not already done and get its OpenAI file id
    file_ids = set()
    retval = []
    for filename, content_filename, description in files:
      if content_filename is None:
        content_filename = filename
      with open(content_filename, 'rb') as f:
        content = f.read().decode('utf-8')
        
      file_id, error = self._add_file(filename, content, description)
      if error:
        errors.append(error)
      else:
        file_ids.add(file_id)
        retval.append((file_id, filename))

    # Disassociate any old files no longer in the context; could ignore errors
    # here but for now complain if unneeded associations remain
    old_context_file_ids.difference_update(file_ids)
    for file_id in old_context_file_ids:
      error = self._disassociate_file(file_id)
      if error:
        errors.append(error)

    # Associated the new files with the assistant
    for file_id in file_ids:
      error = self._associate_file(file_id)
      if error:
        errors.append(error)

    self.__purge_uploaded_files()

    if errors:
      error = '\n'.join(errors)
    else:
      error = None
      
    return retval, error

  def get_filename(self, file_id):
    """Get the filename for the given Open AI file id.  Returns None if the file_id
    is no longer in the files being managed by this assistant."""

    if file_id not in self.file_id_to_filename:
      return None

    return self.file_id_to_filename[file_id][0]

  def get_uploaded_files(self):
    """Get a list of (filename, file_id) for all the files that have been uploaded into
    the user's account. This includes files in the current context, and also any others
    that were previously uploaded but not yet removed."""

    retval = []
    for filename, content_hash in self.uploaded_files:
      retval.append((filename, self.uploaded_files[(filename, content_hash)][0]))

    return retval

  def create_thread(self):
    """Create a new chat thread managed by this assistant.  Returns (thread, error)
    where thread is a Thread object on success and error is an error string on failure."""

    try:
      thread = Thread(self)
    except:
      error = []
      error.append("Failed to create chat thread")
      error.append(traceback.format_exc())
      return None, '\n'.join(error)

    self.threads[thread.get_id()] = thread

    return thread, None

  @classmethod
  def delete_uploaded_file_from_account(cls, file_id, rate=None):
    """Delete the file identified by the given OpenAI file id.  This is a class method
    so that it can be used to delete stale file ids after the assistant instance is gone.
    Returns an error if the file was not deleted, or otherwise None."""

    import openai
    
    client = _GetClient()
    if rate is None:
      rate = RateLimit()

    try:
      rate.tick()
      response = client.files.delete(
        file_id=file_id
      )

      if not response.deleted:
        return f"Failed to delete file {file_id} without exception"
      else:
        return None

    except openai.NotFoundError:
      pass

    except:
      return f"Failed to delete file {file_id}:\n" + traceback.format_exc()

    return None

  @classmethod
  def delete_vector_store_from_account(cls, vector_store_id, rate=None):
    """Delete the given vector_store by ID from the OpenAI account. Returns an error if the
    vector store was not deleted, or otherwise None."""

    import openai
    
    client = _GetClient()
    if rate is None:
      rate = RateLimit()
    
    try:
      rate.tick()
      vector_store_files = client.vector_stores.files.list(vector_store_id).data
    except openai.NotFoundError:
      return None
    except:
      return traceback.format_exc()

    for vector_store_file in vector_store_files:
      cls.delete_uploaded_file_from_account(vector_store_file.id, rate)
    
    try:
      rate.tick()
      vector_store = client.vector_stores.delete(vector_store_id)
      if not vector_store.deleted:
        return "Vector store deletion failed without exception"
    except openai.NotFoundError:
      pass
    except:
      return f"Failed to delete vector store {vector_store_id}:\n" + traceback.format_exc()

    return None

  @classmethod
  def delete_assistant_from_account(cls, assistant_id, rate=None):
    """Delete the given assistant by ID from the OpenAI account. Handles both
    legacy assistant IDs (asst_*) via the beta API and new vector store IDs (vs_*).
    Returns an error if the resource was not deleted, or otherwise None."""

    import openai

    client = _GetClient()
    if rate is None:
      rate = RateLimit()

    # New-style IDs are vector store IDs
    if assistant_id.startswith('vs_'):
      return cls.delete_vector_store_from_account(assistant_id, rate)

    # Legacy assistant IDs — use beta API which may be removed
    try:
      rate.tick()
      assistant = client.beta.assistants.retrieve(assistant_id)
      if not assistant:
        return None
    except openai.NotFoundError:
      return None
    # This occurs with old assistant IDS not starting with asst-
    except openai.BadRequestError:
      if not assistant_id.startswith('asst_'):
        return None
      return "Assistant lookup failed during deletion (BadRequestError):\n" + traceback.format_exc()
    except AttributeError:
      # beta API may have been removed from the openai package
      return None
    except:
      return "Assistant lookup failed during deletion (other):\n" + traceback.format_exc()
    
    # This also removes files
    if assistant.tool_resources:
      for vector_store_id in assistant.tool_resources.file_search.vector_store_ids:
        cls.delete_vector_store_from_account(vector_store_id, rate)
      
    try:
      rate.tick()
      assistant = client.beta.assistants.delete(assistant_id)
      if not assistant.deleted:
        return "Assistant deletion failed without exception"
    except openai.NotFoundError:
      pass
    except:
      return "Assistant deletion failed:\n" + traceback.format_exc()

    return None

  @classmethod
  def delete_matching_resources(cls, untitled_vector_stores=False):
    """Delete all vector stores and files that were created by this OpenAI integration
    from the OpenAI account.

    When untitled_vector_stores is True this deletes also all untitled vector stores,
    which were created in large numbers by earlier versions of Wing and the OpenAI API.

    This is used for cleanup when normal resource removal fails. Use with caution, since
    this may remove resources that are in use!"""

    client = _GetClient()
    rate = RateLimit(5)

    # Remove all files created by Wing
    try:
      files = client.files.list().data
    except:
      files = []

    for file in files:
      if file.filename.startswith('wing:'):
        print("Deleting file_id", file.id)
        cls.delete_uploaded_file_from_account(file.id, rate)

    # Remove legacy assistants created by Wing (beta API may be removed)
    try:
      after_id = ''
      to_delete = []
      while 1:
        rate.tick()
        assistants = client.beta.assistants.list(limit=100, after=after_id).data

        for assistant_id in to_delete:
          print("Deleting assistant_id", assistant_id)
          cls.delete_assistant_from_account(assistant_id, rate)

        after_id = ''
        to_delete = []
        for assistant in assistants:
          if assistant.description == 'wing:assistant':
            to_delete.append(assistant.id)
          after_id = assistant.id

        if not assistants:
          break
    except:
      # beta assistants API may have been removed
      pass

    # Remove vector_stores created by Wing
    after_id = ''
    to_delete = []
    while 1:
      rate.tick()
      vector_stores = client.vector_stores.list(limit=100, after=after_id).data

      for vector_store_id in to_delete:
        print("Deleting vector_store_id", vector_store_id)
        cls.delete_vector_store_from_account(vector_store_id, rate)
        
      after_id = ''
      to_delete = []
      for vector_store in vector_stores:
        if not vector_store.name and not untitled_vector_stores:
          continue
        if vector_store.name is None or vector_store.name.startswith('wing:'):
          to_delete.append(vector_store.id)
        after_id = vector_store.id
    
      if not vector_stores:
        break
    
  def __clear(self):
    """Clear/initialize this instance to blank"""

    self.filename_to_file_ids = collections.defaultdict(set)
    self.file_id_to_filename = {}
    self.hash_to_file_id = {}
    self.context_files = set()
    self.file_id_descriptions = {}
    self.uploaded_files = {}
    self.threads = {}
    self.instructions = None
    self.model = None

  def _add_file(self, filename, content, description):
    """Add the given file to this assistant.  The file may either be used for
    context or referenced from individual messages in threads.  Files are
    uploaded if not already there and updated if the content changes.

    Returns (file_id, error) where file_id is the Open AI file id on success
    and error is an error string on failure."""

    content_hash = self.__get_file_hash(filename, content)
    
    # If the file has already been added, just make sure it has been uploaded
    if content_hash in self.hash_to_file_id:
      file_id = self.hash_to_file_id[content_hash]
      if description is not None:
        self.file_id_descriptions[file_id] = description
      return self.__ensure_file_upload(filename, content)

    purge_files = filename in self.filename_to_file_ids
    remove_file_id = self.filename_to_file_ids.get(filename, None)
    
    # Make sure the new version of the file is uploaded
    file_id, error = self.__ensure_file_upload(filename, content)
    if error:
      return None, error

    # Add it to local structures
    self.filename_to_file_ids[filename].add(file_id)
    self.file_id_to_filename[file_id] = filename
    self.hash_to_file_id[content_hash] = file_id
    if description is not None:
      self.file_id_descriptions[file_id] = description

    # This file replaced an old content version so that upload may no longer be needed
    if purge_files:
      self.__purge_uploaded_files(remove_file_id)

    return file_id, None

  def __get_file_hash(self, filename, content):
    """Get md5 hash of the given filename and its content, to uniquely identify
    that version of the file"""

    combined = filename.encode('utf-8') + b'\n|||\n' + content.encode('utf-8')
    content_hash = hashlib.md5(combined).hexdigest()

    return content_hash

  def __ensure_file_upload(self, filename, content):
    """Make sure that the given file has been uploaded.  Returns (file_id, error)
    where file_id is the Open AI file id for the file on success and error is an
    error string if the upload failed."""

    content_hash = self.__get_file_hash(filename, content)

    if (filename, content_hash) in self.uploaded_files:
      file_id = self.uploaded_files[(filename, content_hash)][0]
      self.uploaded_files[(filename, content_hash)] = (file_id, time.time())
      return file_id, None

    client = _GetClient()

    file_obj = io.BytesIO(content.encode('utf-8'))

    try:
      file = client.files.create(
        file=('wing:'+filename, file_obj),
        purpose="assistants"
      )
    except:
      exc = traceback.format_exc()
      return None, exc

    self.uploaded_files[(filename, content_hash)] = (file.id, time.time())

    return file.id, None

  def _associate_file(self, file_id):
    """Associate the given file ID with this assistant, so it is considered to
    be part of the context for all threads managed by it.  Returns an error string
    on failure or None on success."""

    if file_id in self.context_files:
      return None

    client = _GetClient()

    try:
      client.vector_stores.files.create(
        vector_store_id=self.vector_store.id,
        file_id=file_id
      )
    except:
      exc = traceback.format_exc()
      return exc

    self.context_files.add(file_id)

    return None

  def _disassociate_file(self, file_id):
    """Remove the given file ID from the files that are associated with this
    assistant, as context for all threads managed by it.  Returns an error string
    on failure or None on success."""

    if file_id not in self.context_files:
      return None

    client = _GetClient()

    try:
      deleted = client.vector_stores.files.delete(
        vector_store_id=self.vector_store.id,
        file_id=file_id
      )
      self.context_files.remove(file_id)
      if not deleted.deleted:
        exc = f"disassociating file {file_id} failed without exception"
      else:
        exc = None
    except:
      exc = traceback.format_exc()

    if exc is None:
      self.__purge_uploaded_files()

    return exc

  def _thread_deleted(self, thread_id):
    """Called by a thread when it is deleted"""

    self.threads.pop(thread_id)
    self.__purge_uploaded_files()

  def __purge_uploaded_files(self, force_id=None):
    """Purge any unneeded uploaded files, based on which files are being used by
    the current assistant context set with set_context() and by messages in the
    threads managed by this assistant."""

    thread_file_ids = set()
    for thread in self.threads.values():
      thread_file_ids.update(thread.file_ids)

    for filename, content_hash in list(self.uploaded_files):
      self.__purge_uploaded_file_from_account(filename, content_hash, thread_file_ids,
                                              force_id)

  def __purge_uploaded_file_from_account(self, filename, content_hash, thread_file_ids,
                                         force_id=None):
    """Purge a single upload if no longer needed. We don't purge files that are in use by
    the assistant context or a thread message and we keep files around for up to six hours
    when unused, as long as the assistant still exists, to avoid having to re-upload them."""

    # Not uploaded or already removed
    if (filename, content_hash) not in self.uploaded_files:
      return
    file_id, upload_time = self.uploaded_files[(filename, content_hash)]

    # Don't remove file if recently uploaded, except if forcing removal of an old version
    # of a file that has been edited because it's unlikely we'll want to use that one
    # again
    if upload_time > time.time() - 6 * 60 * 60 and file_id != force_id:
      return

    # Don't remove files if in current context
    if file_id in self.context_files:
      return

    # Don't remove files referenced from a message in an active thread
    if file_id in thread_file_ids:
      return

    error = self.delete_uploaded_file_from_account(file_id)
    if not error:
      del self.uploaded_files[(filename, content_hash)]

_thread_counter = 0

class Thread:
  """A single chat thread that belongs to an assistant. Conversation history is
  maintained via previous_response_id chaining with the Responses API."""

  def __init__(self, assistant):
    """Create a new thread and associate it with the given Assistant instance.
    No API call is needed — this is a local object."""

    global _thread_counter
    _thread_counter += 1
    self.assistant = assistant
    self.thread_id = f'thread-{_thread_counter}-{int(time.time())}'
    self.file_ids = []
    self.active_run = None
    self.last_response_id = None
    self.pending_prompt = None
    self._last_run = None

  def get_id(self):
    """Get the local thread id for this thread"""

    return self.thread_id

  def add_message(self, message, filename, content, focus):
    """Add a new user message to the thread with the given message. The filename and
    content are the main context for the message, the one currently open in an editor, and
    focus is the selection or other subset of that file that the user has focused on.
    Returns (message_id, description) where message_id is a synthetic message ID or None on
    failure, and description is a verbose description of the data sent to OpenAI for the
    request, appropriate for logging."""

    description = []

    # Add the main context file to the assistant.  If already there it's just reused.
    if filename is not None:
      file_id, error = self.assistant._add_file(filename, content, None)
      if error:
        description.append(f"Failed to upload message file {filename} with error:")
        description.append(error)
      else:
        description.append(f"Uploaded message file {filename} (file id {file_id})")
        self.file_ids.append(file_id)
    else:
      file_id = None

    # Assemble the prompt to write the user's request in, to make it clear to the AI
    # what this message is.
    if file_id is None:
      preamble = 'The user currently has no files open.'

    else:
      self.assistant._associate_file(file_id)

      file_open = f'The user currently has the file `{filename}` open'

      if not focus:
        preamble = file_open + " and is asking about all the code in that file."

      elif focus:
        preamble = file_open + f''' and is asking about the following subset of code in that file:

  ```
  {focus}
  ```
  '''
      else:
        preamble = file_open + " but has nothing selected and may be asking a general question or referring to the file as a whole."

    # Include a list of filenames for context so the model knows what files are
    # available via file_search
    file_id_list = []
    if file_id is not None:
      file_id_list.append(f'{filename} -- This is the currently open file')

    for context_file_id in self.assistant.context_files:
      context_filename = self.assistant.file_id_to_filename[context_file_id]
      desc = self.assistant.file_id_descriptions.get(context_file_id, "No description available")
      file_id_list.append(f'{context_filename} - {desc}')

    num_files = len(file_id_list)
    file_id_list = '\n'.join(file_id_list)

    if file_id_list:
      preamble += f'''  {num_files} file(s) have been uploaded and are available via file search.  You can use the file_search tool to look at these files.  These are the file names and descriptions for all the uploaded files available to you:

              {file_id_list}'''

    prompt = f'''This automated call is being made by a plugin to Wing IDE that allows you to help the user with their code.  This preamble adds additional information for the user's request given below.  {preamble}

This is the user's message / request:

```
{message}
```

Your reply can be a mix of text and code delimited by triple backticks.
'''

    description.append("Sending AI Chat message:")
    description.append(prompt)

    # Store prompt for use in create_run()
    self.pending_prompt = prompt

    msg_id = f'msg-{_thread_counter}-{int(time.time())}'

    description = '\n'.join(description)

    return msg_id, description

  def create_run(self):
    """Create a run of this thread.  The run blocks until the response is complete.
    Returns (run, error) where run is the Run instance and error is an error string
    if creation of the response failed."""

    # Only allow one thread run at a time
    assert self.active_run is None

    run = Run(self)
    self.active_run = run
    error = run.start()
    self._last_run = run

    return run, error

  def delete(self):
    """Delete this thread.  Just cleans up local state since there are no server-side
    thread resources to remove.  Returns True."""

    thread_id = self.get_id()
    self.assistant._thread_deleted(thread_id)

    return True

  def get_messages(self):
    """Get the response text from the last completed run. Returns (messages, error)
    where messages is a list of (message_id, assistant_id, role, message, file_ids).
    Kept for interface compatibility."""

    run = self._last_run
    if run is None:
      return [], "No run has been completed"

    text = run.get_response_text()
    if text is None:
      return [], "No response available"

    msg_id = run.response.id if run.response else 'unknown'
    retval = [(msg_id, None, 'assistant', text, [])]

    return retval, None

class Run:
  """A run that calls the Responses API.  The call blocks until the response is complete."""

  def __init__(self, thread):
    """Create the run instance for running the given chat thread"""

    self.thread = thread
    self.response = None

  def start(self):
    """Start the run by calling the Responses API.  This blocks until the response
    is complete.  Returns an error string on failure or None on success."""

    client = _GetClient()

    tools = []
    if self.thread.assistant.vector_store is not None:
      tools.append({
        'type': 'file_search',
        'vector_store_ids': [self.thread.assistant.vector_store.id]
      })

    kwargs = {
      'model': self.thread.assistant.model,
      'instructions': self.thread.assistant.instructions,
      'input': self.thread.pending_prompt,
      'tools': tools,
      'store': True,
    }

    if self.thread.last_response_id is not None:
      kwargs['previous_response_id'] = self.thread.last_response_id

    try:
      self.response = client.responses.create(**kwargs)
      self.thread.last_response_id = self.response.id
      self.thread.active_run = None
      error = None
    except:
      error = traceback.format_exc()
      self.thread.active_run = None

    return error

  def cancel(self):
    """Cancel this run.  No-op for the synchronous Responses API — client-side
    cancellation works by ignoring the callback."""

    return None

  def check_status(self):
    """Check the status of this run.  Returns (error, status, completion_time).
    Since the Responses API is synchronous, this returns immediately."""

    if self.response is not None:
      return None, 'completed', time.time()
    else:
      return 'No response available', 'failed', time.time()

  def get_cost(self):
    """Get the total cost of this run.  This is computed based on the number of OpenAI
    tokens consumed in the message and response and the configured costs per 1K tokens.
    Returns (cost, error) where error is an error string if the cost could not be
    obtained."""

    if self.response is None:
      return 0.0, "No response available"

    if self.response.usage is None:
      return 0.0, None

    input_tokens = self.response.usage.input_tokens
    output_tokens = self.response.usage.output_tokens
    cost = 0.0
    cost += (input_tokens / 1000.) * _GetConfig('prompt_tok')
    cost += (output_tokens / 1000.) * _GetConfig('result_tok')

    return cost, None

  def get_response_text(self):
    """Get the text of the response.  Returns None if no response is available."""

    if self.response is None:
      return None

    return self.response.output_text

def ParseMessage(message):
  """Parse the given output message from the AI to remove any code block indicators and
  generate a mapping from line number in the converted message to mime type for that line,
  if it was part of a code block. No entry in the map means that line is plain text.
  Returns (converted_output, mime_map)."""

  codemap = {}
  filtered = []

  in_code = None
  skipped = 0
  for i, line in enumerate(message.splitlines(keepends=True)):
    if line.lstrip().startswith('```'):
      if in_code is not None:
        in_code = None
      else:
        in_code = line[len('```'):].strip().lower()
        if not in_code:
          in_code = 'python'
        in_code = __GetCodeMimeType(in_code)
      skipped += 1
    elif in_code:
      codemap[i-skipped] = in_code
      filtered.append(line)
    elif not in_code:
      filtered.append(line)

  return ''.join(filtered), codemap

_kIndicatorMimeTypes = {
  "python": "text/x-python",
  "javascript": "application/javascript",
  "json": "application/json",
  "html": "text/html",
  "css": "text/css",
  "xml": "application/xml",
  "markdown": "text/markdown",
  "bash": "application/x-sh",
  "yaml": "text/yaml",
  "sql": "application/sql",
  "plaintext": "text/plain",
  "java": "text/x-java-source",
  "csharp": "text/x-csharp",
  "cpp": "text/x-c++src",
  "c": "text/x-csrc",
  "php": "application/x-httpd-php",
  "ruby": "text/x-ruby",
  "perl": "text/x-perl",
  "go": "text/x-go",
  "swift": "text/x-swift",
  "typescript": "video/mp2t",
  "powershell": "application/x-powershell",
  "rust": "text/x-rustsrc",
  "dart": "application/vnd.dart",
  "kotlin": "text/x-kotlin",
  "scala": "text/x-scala",
  "elixir": "text/x-elixir",
  "shell": "application/x-sh",
  "lua": "text/x-lua",
  "groovy": "text/x-groovy",
  "r": "text/x-rsrc",
  "matlab": "text/x-matlab",
  "julia": "text/x-julia",
  "fortran": "text/x-fortran",
  "pascal": "text/x-pascal",
  "haskell": "text/x-haskell",
  "erlang": "text/x-erlang",
  "clojure": "text/x-clojure"
}

def __GetCodeMimeType(block_type):
  """Utility to get the mime type for the given OpenAI block type.  Note that OpenAI
  block types are not well documented and the mapping was produced by asking OpenAI
  to create it."""

  block_type = block_type.lower().strip()
  if block_type == '':
    return 'text/x-python'
  else:
    return _kIndicatorMimeTypes.get(block_type, 'text/plain')
    

class AssistantServer:
  
  def __init__(self):
    
    self._assistant = Assistant()
    
    self.host = '127.0.0.1'
    
    self.client_socket = None
    
    self._socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    self._socket.setblocking(True)
    self._socket.bind((self.host, 0))
    self._socket.listen()
    
    self.port = self._socket.getsockname()[1]

    self._message_length = None
    self._pending_data = b''

  def _error(self, error, exc):
    
    # Utility on the theory that trying print when sys.stderr has been closed
    # might be worth something in some cases during dev
    def try_write(txt):
      try:
        sys.stderr.writable(txt)
      except:
        print(txt)
        
    if exc is not None:
      try_write(error+':\n')
      try_write(exc+'\n')
    else:
      try_write(error+'\n')
    
  def _handle_request(self):

    # Reading the 4-byte header that gives the message length
    if self._message_length is None:
      
      # Read up to what's missing from the 4 bytes of message length
      data = self.client_socket.recv(4-len(self._pending_data))
      if not data:
        self._error("Connection closed unexpectedly while reading message header", None)
        return False
      self._pending_data += data
      
      # Exit if we still don't have all the 4 bytes      
      if len(self._pending_data) < 4:
        return True

      # Unpack the message length
      try:
        message_length = struct.unpack('>I', self._pending_data)[0]
      except:
        self._error("JSON unpack of message header failed", traceback.format_exc())
        return False
      self._pending_data = b''
      self._message_length = message_length
      
    # Read up to what's missing from the message
    chunk = self.client_socket.recv(self._message_length - len(self._pending_data))
    if not chunk:
      self._error("Connection closed unexpectedly", traceback.format_exc())
      return False
    self._pending_data += chunk

    # Exit if we don't yet have the full message
    if len(self._pending_data) < self._message_length:
      return True

    # Unpack the response
    try:
      response_data = json.loads(self._pending_data.decode('utf-8'))
    except json.JSONDecodeError:
      self._error("JSON unpack of message body failed", traceback.format_exc())
      return False
    else:
      self._pending_data = b''
      self._message_length = None

    # Execute the request in a separate thread
    method = response_data['method']
    args = response_data['args']
    kwargs = response_data['kwargs']
    callback_id = response_data['callback_id']

    # Note that _execute_method will write the response on the socket while the
    # main thread is blocked on reading it.  This is supposedly OK to do, although
    # sources predictably vary on this since synchronization is harder to get right.
    # In our very simple case it's probably fine, though.
    execute_thread = threading.Thread(target=self._execute_method,
                                      args=(method, args, kwargs, callback_id))
    execute_thread.start()
    
    return method != 'delete'

  def _execute_method(self, method_name, args, kwargs, callback_id):
    method = getattr(self, method_name)
    result = method(*args, **kwargs)
    response = json.dumps({'result': result, 'callback_id': callback_id})
    message = response.encode('utf-8')
    try:
      self.client_socket.sendall(struct.pack('>I', len(message)) + message)
    except:
      if method_name != 'delete':
        raise
      else:
        print("Assistant deleted")
        sys.exit(0)

  def set_config(self, config):
    _OverrideConfig(config)
    return ()
    
  def start(self):
    self.client_socket, _ = self._socket.accept()
    self.client_socket.setblocking(True)
    
    while True:
      if not self._handle_request():
        break
      
    print("Exiting server")
    
  def create(self, name):
    return self._assistant.create(name)
      
  def create_thread(self):
    
    thread, error = self._assistant.create_thread()
    if thread is not None:
      thread_id = thread.get_id()
    else:
      thread_id = None
      
    return thread_id, error

  def delete_thread(self, thread_id):
    
    thread = self._assistant.threads.get(thread_id, None)
    if thread is None:
      return (True, )
    
    deleted = thread.delete()
    return (deleted, )
    
  def set_context(self, files):
    
    file_ids, error = self._assistant.set_context(files)
    uploaded_files = self._assistant.get_uploaded_files()
    return file_ids, uploaded_files, error
  
  def run_request(self, thread_id, request, filename, content_file, selected_text):
    
    errors = []
    
    with open(content_file, 'rb') as f:
      content = f.read().decode('utf-8')
      
    thread = self._assistant.threads[thread_id]
    message_id, description = thread.add_message(request, filename, content, selected_text)
    if message_id is None:
      errors.append("Could not add message to chat thread:")
      errors.append(description)
      return 'failed', None, {} , 0.0, self._assistant.get_uploaded_files(), '\n'.join(errors)
    
    errors = [description]

    # create_run() blocks until the response is complete
    run, error = thread.create_run()
    if error:
      errors.append("Could not create run:")
      errors.append(error)
      return 'failed', None, {} , 0.0, self._assistant.get_uploaded_files(), '\n'.join(errors)

    error, status, end_time = run.check_status()
    if error:
      errors.append("Run failed:")
      errors.append(error)
      return 'failed', None, {} , 0.0, self._assistant.get_uploaded_files(), '\n'.join(errors)

    cur_cost, error = run.get_cost()
    if error:
      errors.append("Could not get cost of run:")
      errors.append(error)
      cur_cost = 0.0

    result = run.get_response_text()
    if result is None:
      errors.append("No response text available")
      parsed_text = None
      mime_types = {}
    else:
      errors.append("AI Chat response:"+'\n'+result)
      parsed_text, mime_types = ParseMessage(result)
    
    return status, parsed_text, mime_types, cur_cost, self._assistant.get_uploaded_files(), '\n'.join(errors)

  def cancel_request(self, thread_id):
    
    thread = self._assistant.threads.get(thread_id, None)
    if thread is None or thread.active_run is None:
      return (None, )
    
    return (thread.active_run.cancel(), )
    
  def delete(self):
    errors = self._assistant.delete()
    return (errors, )
  
def StartServer(result_cb, output_cb):
  """Start the OpenAI Assistant server. Calls result_cb(exit_code, error, host, port) with
  an exit_code != -1 if the server processed exited, error set to a string suitable for
  logging if there are errors, and host and port set once the server has started and is
  listening for connections. Calls output_cb(output) when there is any output from the
  server that is suitable for logging."""
  
  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("Starting AI chat service process", 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, None, None)
  
  all_output = []
  def output_received(obj, output):
    
    if output:
      output_cb(output)
      
    all_output.append(output)
    
    txt = ''.join(all_output)
    if 'STARTED=' in txt and '\n' in txt:
      hostport = txt[txt.find('STARTED=')+len('STARTED='):]
      hostport = hostport[:hostport.find('\n')]
      host, port = hostport.split(':')
      port = int(port.strip())
      result_cb(-1, None, host, port)
      
  def error_output_received(obj, output):
    if output:
      output_cb(output)

  def terminated(obj):
    exit_code = proc.GetExitCode()
    error = proc.GetStdErrOutput()
    result_cb(exit_code, error, None, None)

  proc.connect('terminated', terminated)
  proc.connect('start-failed', start_failed)
  proc.connect('received-output', output_received)
  proc.connect('received-stderr-output', error_output_received)

  proc.Start()

  description = [
      spawn._win32_cmd_line_from_argv(args),
        'WINGHOME={}'.format(_winghome.kWingHome),
        'WINGSETTINGSDIR={}'.format(_winghome.kUserSettings),
    ]
  description = '\n'.join(description)

  return proc, description
  
  
class AssistantClient:
  """Client implementation for connecting to and sending requests to AssistantServer.
  This class just contains the support code for communication and not the public
  methods used by the IDE."""
  
  def __init__(self):
    self._handler = None
    self._socket = None
    self._callback_map = {}
    self._ignore_callbacks = set()
    self._message_length = None
    self._pending_data = b''
    self._pending_requests = []
  
  def connect_to_server(self, host, port):
    self._socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    self._socket.connect((host, port))
    self._socket.setblocking(False)
    
    # Assumes pending messages don't exceed buffer size, which they shouldn't
    while self._pending_requests:
      callback_id, header, message = self._pending_requests.pop(0)
      self._socket.sendall(header + message)
      
  def disconnect(self):
    assert not self._callback_map
    if self._socket is not None:
      self._socket.close()
    self._socket = None
    self._pending_data = b''
    
  def _ensure_handler(self):
    
    from guiutils import wgtk
    
    if not self._callback_map and self._handler is not None:
      wgtk.idle_remove(self._handler)

    elif self._callback_map and self._handler is None:
      self._handler = wgtk.idle_add(self._handle_response)

  def _handle_response(self):

    # Exit if not waiting for a response and remove the idle
    if not self._callback_map or self._socket is None:
      self._handler = None
      return False
    
    # Exit if nothing to read yet, but keep the idle active
    readable, _, _ = select.select([self._socket], [], [], 0)
    if not readable:
      return True

    # Reading the 4-byte header that gives the message length
    if self._message_length is None:
      
      # Read up to what's missing from the 4 bytes of message length
      data = self._socket.recv(4-len(self._pending_data))
      if not data:
        self._error("Connection closed unexpectedly")
        self._handler = None
        return False
      self._pending_data += data
      
      # Exit if we still don't have all the 4 bytes      
      if len(self._pending_data) < 4:
        return True

      # Unpack the message length
      try:
        message_length = struct.unpack('>I', self._pending_data)[0]
      except:
        self._error("JSON unpack of message header failed")
        self._handler = None
        return False
      self._pending_data = b''
      self._message_length = message_length
      
    # Read up to what's missing from the message
    chunk = self._socket.recv(message_length - len(self._pending_data))
    if not chunk:
      self._error("Connection closed unexpectedly")
      self._handler = None
      return False
    self._pending_data += chunk

    # Exit if we don't yet have the full message
    if len(self._pending_data) < self._message_length:
      return True

    # Unpack the response
    try:
      response_data = json.loads(self._pending_data.decode('utf-8'))
    except json.JSONDecodeError:
      self._error("JSON unpack of message body failed")
      self._handler = None
      return False
    else:
      self._pending_data = b''
      self._message_length = None

    # Call the result callback
    callback_id = response_data.get('callback_id')
    if callback_id in self._ignore_callbacks:
      self._ignore_callbacks.remove(callback_id)
      return True
    self._callback_map[callback_id](*response_data.get('result'))
    del self._callback_map[callback_id]
    
    return True

  def _send_request(self, method, callback, *args, **kwargs):
    
    # Add callback to map
    callback_id = id(callback)
    self._callback_map[callback_id] = callback
    self._ensure_handler()
    
    request = json.dumps({'method': method, 'args': args, 'kwargs': kwargs,
                          'callback_id': callback_id})
    message = request.encode('utf-8')
    header = struct.pack('>I', len(message))
    
    # Early call during connection
    if self._socket is None:
      self._pending_requests.append((callback_id, header, message))
      return callback_id

    self._socket.sendall(header + message)
    
    return callback_id

  def close(self):
    self._socket.close()

  def _error(self, error):
    raise NotImplementedError("subclass responsibility")

def ChatAPIMain():
  """This implements the server used for calling OpenAI Responses API for AI Chat. It
  starts the server and when ready prints a status message that indicates what the
  connection host and port should be for the client."""

  server = AssistantServer()
  print(f"STARTED={server.host}:{server.port}\n")
  sys.stdout.flush()
  server.start()


#########################################################################
if __name__ == '__main__':
  #try:
    #import _wingdbstub
  #except ImportError:
    #pass
  ChatAPIMain()


