import os
from openai import OpenAI

client = OpenAI(
    api_key=os.environ["RADIUM_API_KEY"],
    base_url="https://api.radium.cloud/v1",
)


class Conversation:
    def __init__(self, model="clarke-1.0", max_history=10, summary_trigger=8):
        self.model = model
        self.messages = [
            {"role": "system", "content": "You are a helpful assistant."}
        ]
        self.max_history = max_history      # hard cap on turns kept
        self.summary_trigger = summary_trigger  # when to summarize old turns

    def ask(self, user_message: str) -> str:
        # 1. Add user message
        self.messages.append({"role": "user", "content": user_message})

        # 2. Summarize if history is getting long
        if len(self.messages) > self.summary_trigger + 1:
            self._summarize()

        # 3. Trim if still over limit
        self._trim()

        # 4. Send to model
        response = client.chat.completions.create(
            model=self.model,
            messages=self.messages,
            temperature=0.7,
            max_tokens=512,
        )

        assistant_reply = response.choices[0].message.content
        self.messages.append({"role": "assistant", "content": assistant_reply})
        return assistant_reply

    def _trim(self):
        """Keep only the system message + last `max_history` turns."""
        # +2 because system msg + pairs of [user, assistant]
        while len(self.messages) > self.max_history * 2 + 1:
            # Don't remove system message at index 0
            self.messages.pop(1)

    def _summarize(self):
        """Replace old user/assistant turns with a summary via the model."""
        # Everything except system message and the most recent 2 turns
        old_turns = self.messages[1:-4]
        if not old_turns:
            return

        summary_request = [
            {"role": "system", "content": "Summarize the following conversation briefly."},
            {"role": "user", "content": "\n".join(f"{m['role']}: {m['content']}" for m in old_turns)},
        ]
        summary_resp = client.chat.completions.create(
            model=self.model,
            messages=summary_request,
            temperature=0.3,
            max_tokens=256,
        )
        summary = summary_resp.choices[0].message.content

        # Replace old turns with a system-level summary
        self.messages = [
            self.messages[0],  # original system prompt
            {"role": "system", "content": f"Conversation summary so far: {summary}"},
        ] + self.messages[-4:]  # keep last 2 user/assistant pairs


# --- Run it ---
chat = Conversation(model="clarke-1.0")

for msg in [
    "What is machine learning?",
    "Can you give me a concrete example?",
    "Who invented the perceptron?",
    "When was that exactly?",
    "How does a neural network learn?",
    "What is backpropagation?",
    "Tell me about transformers.",
    "Who wrote the Attention Is All You Need paper?",
    "Summarize everything we discussed.",
]:
    print(f"\nUser: {msg}")
    reply = chat.ask(msg)
    print(f"Assistant: {reply}")

print(f"\nFinal conversation has {len(chat.messages)} messages.")