import requests
import json
import time
import sys
import os
from dotenv import load_dotenv
load_dotenv()

# === CONFIGURATION ===
AUTH_TOKEN =  os.getenv("AUTH_TOKEN")  
BASE_URL = "https://chat.cumin.dev/api"
HEADERS = {
    "Authorization": f"Bearer {AUTH_TOKEN}",
    "Content-Type": "application/json"
}

# Fixed initial prompt (R&D assistant, Arabic only)
INITIAL_PROMPT = """You are an expert R&D assistant specialized in AI, Machine Learning, and LLMs. Provide 3-5 creative, underexplored research directions. For each idea: short title, brief explanation, why it's promising (with novelty evidence), and at least one reference to a paper or source. Be concise but insightful. Reply only in Arabic."""

def create_conversation(title="AI R&D Session"):
    url = f"{BASE_URL}/conversations"
    payload = {"title": title}
    resp = requests.post(url, headers=HEADERS, json=payload)
    if resp.status_code in (200, 201):
        conv_id = resp.json().get("id")
        print(f"[✓] New conversation created: {conv_id}")
        return conv_id
    else:
        print(f"[✗] Failed to create conversation: {resp.status_code} - {resp.text}")
        return None

def send_message(conv_id, text):
    url = f"{BASE_URL}/conversations/{conv_id}/chat"
    payload = {"message": text, "files": []}
    resp = requests.post(url, headers=HEADERS, json=payload)
    if resp.status_code == 200:
        return True
    else:
        print(f"[✗] Send failed: {resp.text}")
        return False

def get_conversation(conv_id):
    url = f"{BASE_URL}/conversations/{conv_id}"
    resp = requests.get(url, headers=HEADERS)
    if resp.status_code == 200:
        return resp.json()
    else:
        print(f"[✗] Failed to fetch conversation: {resp.text}")
        return None

def get_last_assistant_reply(conv_data):
    if not conv_data or "messages" not in conv_data:
        return None
    for msg in reversed(conv_data["messages"]):
        if msg.get("role") == "assistant":
            return msg.get("content")
    return None

def wait_for_assistant_reply(conv_id, last_msg_count, timeout=40, poll_interval=2):
    start = time.time()
    while time.time() - start < timeout:
        conv = get_conversation(conv_id)
        if conv and "messages" in conv:
            if len(conv["messages"]) > last_msg_count:
                new_msg = conv["messages"][-1]
                if new_msg["role"] == "assistant":
                    return new_msg["content"]
        time.sleep(poll_interval)
    print("[!] Timeout waiting for assistant reply.")
    return None

def initialize_conversation(conv_id):
    """Send the initial R&D prompt and return the assistant's first reply."""
    print("\n[→] Sending initial R&D prompt (Arabic only)...")
    if not send_message(conv_id, INITIAL_PROMPT):
        return None
    print("[→] Waiting for assistant to generate ideas...")
    conv_data = get_conversation(conv_id)
    if not conv_data:
        return None
    msg_count = len(conv_data.get("messages", []))
    reply = wait_for_assistant_reply(conv_id, msg_count)
    return reply

# === MAIN LOOP ===
def main():
    print("\n" + "="*60)
    print("R&D ASSISTANT (AI/ML/LLM) – Arabic responses only")
    print("Commands: /new  -> start a fresh conversation (with initial prompt)")
    print("          /exit -> quit")
    print("="*60 + "\n")

    conv_id = None
    choice = input("Start a new conversation? (y/n): ").strip().lower()
    if choice == 'y':
        conv_id = create_conversation()
        if not conv_id:
            print("Could not create conversation. Exiting.")
            sys.exit(1)
        # Send the initial Arabic R&D prompt automatically
        first_reply = initialize_conversation(conv_id)
        if first_reply:
            print("\n" + "="*60)
            print("ASSISTANT'S INITIAL RESPONSE (Arabic):")
            print("="*60)
            print(first_reply)
        else:
            print("[!] Could not get initial response. You can still chat.")
    else:
        conv_id = input("Paste an existing conversation ID: ").strip()
        print(f"[i] Using existing conversation: {conv_id}")

    # Interactive loop (continue conversation)
    while True:
        user_input = input("\nYou: ").strip()
        if not user_input:
            continue

        if user_input == "/exit":
            print("Goodbye!")
            break

        if user_input == "/new":
            print("\n[!] Creating a fresh conversation...")
            new_id = create_conversation()
            if new_id:
                conv_id = new_id
                # Send initial prompt for the new conversation too
                first_reply = initialize_conversation(conv_id)
                if first_reply:
                    print("\n" + "="*60)
                    print("ASSISTANT'S INITIAL RESPONSE (Arabic):")
                    print("="*60)
                    print(first_reply)
                else:
                    print("[!] No initial response. Continue chatting.")
            else:
                print("[✗] Failed to create new conversation. Staying in current.")
            continue

        # Normal message
        if not send_message(conv_id, user_input):
            continue

        # Wait for reply
        conv_data = get_conversation(conv_id)
        if not conv_data:
            continue
        msg_count = len(conv_data.get("messages", []))
        print("Assistant is thinking...", end="", flush=True)
        reply = wait_for_assistant_reply(conv_id, msg_count)
        print("\r" + " " * 30 + "\r", end="")
        if reply:
            print(f"\nAssistant (Arabic): {reply}")
        else:
            print("\nAssistant: (No response)")

if __name__ == "__main__":
    main()