import sys
import os
import json
import subprocess
from argparse import ArgumentParser

from anthropic import Anthropic, AI_PROMPT, HUMAN_PROMPT
from openai import OpenAI


ASSISTANT_CONFIG_FILE = "/root/tools/assistant.json"
SYSTEM_PROMPT = "You're a software engineer working on a team project. Your job is to write technical summaries of code files that summarize the purpose and key elements of the code. Your summarize will be used by other engineers to quickly understand the code and determine it's use cases and relevance to their work. You should reference major parts of code by their line numbers. You should also include any important classes, functions, or main tasks performed by the code. You should aim for simplicity and clarity in your summary."
INSTANCE_PROMPT = lambda code: f"Here's the code you need to summarize:\n\n{code}\n\nPlease respond with a summary of the code below:\n"
LOG_FILE = '/root/.assistant_data'


def call_anthropic(cfg, code, model_name, temperature, top_p):
    # Create messages
    message = f"{HUMAN_PROMPT}{SYSTEM_PROMPT}\n{HUMAN_PROMPT}{code}\n{AI_PROMPT}"

    # Create Anthropic instance
    client = Anthropic(api_key=cfg["ANTHROPIC_API_KEY"])
    completion = client.completions.create(
        model=model_name,
        prompt=message,
        max_tokens_to_sample=6_0000,
        temperature=temperature,
        top_p=top_p,
    )
    response = completion.completion
    input_tokens = client.count_tokens(message)
    output_tokens = client.count_tokens(response)
    with open(LOG_FILE, '+a') as f:
        print(json.dumps({
            'model_name': model_name,
            'temperature': temperature,
            'top_p': top_p,
            'input_tokens': input_tokens,
            'output_tokens': output_tokens,
        }), file=f)
    return response


def call_openai(cfg, code, model_name, temperature, top_p):
    # Create messages
    messages = [
        {
            "role": "system",
            "content": SYSTEM_PROMPT,
        },
        {"role": "user", "content": INSTANCE_PROMPT(code)},
    ]
    # Create OpenAI instance
    client = OpenAI(api_key=cfg["OPENAI_API_KEY"])
    response = client.chat.completions.create(
        messages=messages,
        model=model_name,
        temperature=temperature,
        top_p=top_p,
    )
    input_tokens = response.usage.prompt_tokens
    output_tokens = response.usage.completion_tokens
    with open(LOG_FILE, '+a') as f:
        print(json.dumps({
            'model_name': model_name,
            'temperature': temperature,
            'top_p': top_p,
            'input_tokens': input_tokens,
            'output_tokens': output_tokens,
        }), file=f)
    completion = response.choices[0].message.content
    return completion


def main(file_path):
    # Read config file
    if not os.path.exists(ASSISTANT_CONFIG_FILE):
        print(f"Summarize disabled. No config file found")
        sys.exit(1)
    cfg = json.load(open(ASSISTANT_CONFIG_FILE))
    # Read the content of the specified file with line numbers
    with open(file_path, "r") as file:
        code = subprocess.run(f'cat {file_path} | grep -n $', shell=True, text=True, capture_output=True).stdout
    model_name = cfg.get("model_name", None)
    if model_name is None:
        print("Summarize disabled. No model name found")
        sys.exit(1)
    temperature = cfg["temperature"]
    top_p = cfg["top_p"]
    if model_name.lower().startswith("claude"):
        summary = call_anthropic(cfg, code, model_name, temperature, top_p)
    elif model_name.lower().startswith("gpt"):
        summary = call_openai(cfg, code, model_name, temperature, top_p)
    else:
        raise ValueError(f"Invalid model name: {model_name}")
    print(f'File: {file_path}')
    print(f'Summary: {summary}')


if __name__ == "__main__":
    # Check if the correct number of command-line arguments is provided
    parser = ArgumentParser()
    parser.add_argument("file_path", help="Path to the file to summarize")
    args = parser.parse_args()
    main(**vars(args))
