import os
from gam import (
    MemoryAgent,
    ResearchAgent,
    OpenAIGenerator,
    OpenAIGeneratorConfig,
    InMemoryMemoryStore,
    InMemoryPageStore,
    DenseRetrieverConfig,
    DenseRetriever,
    IndexRetrieverConfig,
    IndexRetriever,
    BM25RetrieverConfig,
    BM25Retriever
)

# 1. Configure and create generator
gen_config = OpenAIGeneratorConfig(
    model_name="gpt-4o-mini",
    api_key=os.getenv("OPENAI_API_KEY"),
    base_url="https://api.openai.com/v1",
    temperature=0.3,
    max_tokens = 256
)
generator = OpenAIGenerator.from_config(gen_config)

# 2. Create memory and page stores
memory_store = InMemoryMemoryStore()
page_store = InMemoryPageStore()

# 3. Create MemoryAgent
memory_agent = MemoryAgent(
    generator=generator,
    memory_store=memory_store,
    page_store=page_store
)

# 4. Memorize documents
documents = [
    "Artificial Intelligence is a branch of computer science...",
    "Machine Learning is a subset of AI...",
    "Deep Learning uses neural networks..."
]

for doc in documents:
    memory_agent.memorize(doc)

# 5. Get memory state
memory_state = memory_store.load()
print(f"Built {len(memory_state.abstracts)} memory abstracts")

# 6. Create ResearchAgent for Q&A

retrievers={}
index_dir = './tmp'
try:
    page_index_dir = os.path.join(index_dir, "page_index")
    if os.path.exists(page_index_dir):
        import shutil
        shutil.rmtree(page_index_dir)
    
    index_config = IndexRetrieverConfig(
        index_dir=page_index_dir
    )
    index_retriever = IndexRetriever(index_config.__dict__)
    index_retriever.build(page_store)
    retrievers["page_index"] = index_retriever
except Exception as e:
    print(f"[WARN] page retriever error: {e}")

try:
    bm25_index_dir = os.path.join(index_dir, "bm25_index")
    if os.path.exists(bm25_index_dir):
        import shutil
        shutil.rmtree(bm25_index_dir)
    
    bm25_config = BM25RetrieverConfig(
        index_dir=bm25_index_dir,
        threads=1
    )
    bm25_retriever = BM25Retriever(bm25_config.__dict__)
    bm25_retriever.build(page_store)
    retrievers["keyword"] = bm25_retriever
except Exception as e:
    print(f"[WARN] BM25 retriever error: {e}")

try:
    dense_index_dir = os.path.join(index_dir, "dense_index")
    if os.path.exists(dense_index_dir):
        import shutil
        shutil.rmtree(dense_index_dir)
    
    dense_config = DenseRetrieverConfig(
        index_dir=dense_index_dir,
        model_name="BAAI/bge-m3"
    )
    dense_retriever = DenseRetriever(dense_config.__dict__)
    dense_retriever.build(page_store)
    retrievers["vector"] = dense_retriever
except Exception as e:
    print(f"[WARN] Dense retriever error: {e}")

research_agent_kwargs = {
    "page_store": page_store,
    "memory_store": memory_store,
    "retrievers": retrievers,
    "generator": generator,
    "max_iters": 5
}

research_agent = ResearchAgent(**research_agent_kwargs)

# 7. Perform research
research_result = research_agent.research(
    request="What is the difference between ML and DL?"
)

research_summary = research_result.integrated_memory


print(f"[OK] Research completed! Iteration count: {len(research_result.raw_memory.get('iterations', []))}")
print(f"Research Summary: {research_summary}")