#!/usr/bin/env python3

import argparse
import json
import os
import sys
import urllib.parse
from configparser import ConfigParser
from typing import Any, Dict

from litellm import completion  # type: ignore[import-not-found]

import zulip


def format_conversation(result: Dict[str, Any]) -> str:
    # Note: Including timestamps seems to have no impact; including reactions
    # makes the results worse.
    zulip_messages = result["messages"]
    if len(zulip_messages) == 0:
        print("No messages in conversation to summarize")
        sys.exit(0)

    zulip_messages_list = [
        {"sender": message["sender_full_name"], "content": message["content"]}
        for message in zulip_messages
    ]
    return json.dumps(zulip_messages_list)


def make_message(content: str, role: str = "user") -> Dict[str, str]:
    return {"content": content, "role": role}


def get_max_summary_length(conversation_length: int) -> int:
    return min(6, 4 + int((conversation_length - 10) / 10))


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--url",
        type=str,
        help="The URL to fetch content from",
        default="https://chat.zulip.org/#narrow/channel/101-design/topic/buddy.20list.20style.20switcher",
    )
    parser.add_argument(
        "--model",
        type=str,
        help="The model name to use for summarization",
        default="huggingface/meta-llama/Llama-3.1-70B-Instruct",
    )
    parser.add_argument(
        "--max-tokens",
        type=int,
        help="The maximum tokens permitted in the response",
        default=300,
    )
    parser.add_argument(
        "--max-messages",
        type=int,
        help="The maximum number of messages fetched from the server",
        default=100,
    )
    parser.add_argument(
        "--verbose",
        type=bool,
        help="Print verbose debugging output",
        default=False,
    )
    args = parser.parse_args()

    config_file = zulip.get_default_config_filename()
    if not config_file:
        print("Could not find the Zulip configuration file. Please read the provided README.")
        sys.exit()

    client = zulip.Client(config_file=config_file)

    config = ConfigParser()
    # Make config parser case sensitive otherwise API keys will be lowercased
    # which is not supported by litellm.
    # https://docs.python.org/3/library/configparser.html#configparser.ConfigParser.optionxform
    config.optionxform = str  # type: ignore[assignment, method-assign]

    with open(config_file) as f:
        config.read_file(f, config_file)

    # Set all the keys in `litellm` as environment variables.
    for key in config["litellm"]:
        if args.verbose:
            print("Setting key:", key)
        os.environ[key] = config["litellm"][key]

    url = args.url
    model = args.model

    base_url, narrow_hash = url.split("#")
    narrow_hash_terms = narrow_hash.split("/")
    channel = int(narrow_hash_terms[2].split("-")[0])
    topic = narrow_hash_terms[4]
    topic = urllib.parse.unquote(topic.replace(".", "%"))

    narrow = [
        {"operator": "channel", "operand": channel},
        {"operator": "topic", "operand": topic},
    ]

    request = {
        "anchor": "newest",
        "num_before": args.max_messages,
        "num_after": 0,
        "narrow": narrow,
        # Fetch raw Markdown, not HTML
        "apply_markdown": False,
    }
    result = client.get_messages(request)
    if result["result"] == "error":
        print("Failed fetching message history", result)
        sys.exit(1)

    conversation_length = len(result["messages"])
    max_summary_length = get_max_summary_length(conversation_length)

    print("Conversation URL:", url)
    print(f"Max summary length: {max_summary_length}")

    intro = f"The following is a chat conversation in the Zulip team chat app. channel: {channel}, topic: {topic}"
    formatted_conversation = format_conversation(result)
    prompt = f"Succinctly summarize this conversation based only on the information provided, in up to {max_summary_length} sentences, for someone who is familiar with the context. Mention key conclusions and actions, if any. Refer to specific people as appropriate. Don't use an intro phrase."
    messages = [
        make_message(intro, "system"),
        make_message(formatted_conversation),
        make_message(prompt),
    ]

    # Send formatted messages to the LLM model for summarization
    response = completion(
        max_tokens=args.max_tokens,
        model=model,
        messages=messages,
    )

    print(
        f"Used {response['usage']['completion_tokens']} completion tokens to summarize {conversation_length} Zulip messages ({response['usage']['prompt_tokens']} prompt tokens)."
    )
    print()
    print(response["choices"][0]["message"]["content"])
