# Requires: Python 3.10+, faiss-cpu==1.7.4, openai==1.10.0, numpy==1.24.4
import faiss
import numpy as np
import openai
import json
import os

openai.api_key = os.getenv("OPENAI_API_KEY")
VECTOR_DIM = 1536  # Matches text-embedding-ada-002

class VectorMemory:
    def __init__(self):
        self.index = faiss.IndexFlatL2(VECTOR_DIM)
        self.memory = []

    def insert(self, text, embedding):
        self.memory.append({'text': text, 'embedding': embedding})
        self.index.add(np.array([embedding]).astype(np.float32))

    def recall(self, query_embedding, k=3):
        if len(self.memory) == 0:
            return []
        D, I = self.index.search(np.array([query_embedding]).astype(np.float32), min(k, len(self.memory)))
        return [self.memory[i]['text'] for i in I[0] if i < len(self.memory)]

class StructuredMemory:
    def __init__(self):
        self.state = {}

    def insert(self, key, value):
        self.state[key] = value

    def recall(self, key):
        return self.state.get(key)

    def update(self, key, value):
        self.state[key] = value

    def to_json(self):
        return json.dumps(self.state, indent=2)

def get_embedding(text):
    try:
        resp = openai.embeddings.create(
            model="text-embedding-ada-002",
            input=[text]
        )
        return resp.data[0].embedding
    except Exception as e:
        print("Embedding error:", e)
        return [0.0] * VECTOR_DIM

def summarize(history):
    prompt = (
        "Summarize the following conversation in 1-2 sentences:\n"
        f"{history}\nSummary:"
    )
    try:
        response = openai.chat.completions.create(
            model="gpt-3.5-turbo",
            messages=[
                {"role": "system", "content": "You summarize conversations for an AI agent context."},
                {"role": "user", "content": prompt}
            ],
            max_tokens=50
        )
        return response.choices[0].message.content.strip()
    except Exception as e:
        print("Summarization error:", e)
        return "Could not summarize."

def agent_respond(prompt):
    try:
        response = openai.chat.completions.create(
            model="gpt-3.5-turbo",
            messages=[
                {"role": "system", "content": "You are a helpful assistant. Use the provided task and memory to help the user."},
                {"role": "user", "content": prompt}
            ],
            max_tokens=60
        )
        return response.choices[0].message.content.strip()
    except Exception as e:
        print("Agent error:", e)
        return "I'm unable to reply right now."

# --- Main Loop (real sample data) ---

vector_memory = VectorMemory()
structured_memory = StructuredMemory()
conversation_history = []

sample_inputs = [
    "Hi, I want to plan a trip to Japan.",
    "Book hotel in Kyoto. task: Book hotel in Kyoto",
    "Remind me where I'm staying.",
    "What's my current task?"
]

for user_msg in sample_inputs:
    print("> User:", user_msg)

    user_emb = get_embedding(user_msg)
    vector_memory.insert(user_msg, user_emb)
    conversation_history.append(user_msg)

    if "task:" in user_msg:
        task = user_msg[user_msg.index("task:") + 5:].strip()
        structured_memory.insert("current_task", task)

    recent = vector_memory.recall(user_emb, k=3)
    summary = summarize(" ".join(conversation_history[-10:]))
    current_task = structured_memory.recall("current_task")

    prompt_parts = [
        f"Task Memory: {current_task or 'None'}",
        f"Summary: {summary}",
        f"Recall: {'; '.join(recent)}",
        f"User: {user_msg}"
    ]
    agent_prompt = "\n".join(prompt_parts)

    agent_reply = agent_respond(agent_prompt)
    print("Agent:", agent_reply)
    conversation_history.append(agent_reply)