"""
Realistic end-to-end workflow test for memstate-langchain.

This script simulates a real-world AI agent memory workflow:
1. An agent starts a new project and stores initial decisions
2. The agent saves a rich task summary via server-side intelligence
3. A new agent session retrieves context before starting work
4. The agent searches for specific information
5. The agent updates a decision (creating a new version)
6. Version history is inspected to audit the change
7. Chat message history is tested across sessions
8. Agent tools are exercised end-to-end

Run with:
    MEMSTATE_API_KEY=mst_... python tests/test_workflow.py

Optional (for full LangChain agent test):
    OPENAI_API_KEY=sk_... MEMSTATE_API_KEY=mst_... python tests/test_workflow.py
"""

import os
import sys
import time
import uuid

# Add the parent directory to the path for direct execution
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

API_KEY = os.environ.get("MEMSTATE_API_KEY", "")
OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY", "")
BASE_URL = os.environ.get("MEMSTATE_BASE_URL", "https://api.memstate.ai")

if not API_KEY:
    print("ERROR: MEMSTATE_API_KEY environment variable is required.")
    print("Get your API key at https://memstate.ai/dashboard")
    sys.exit(1)

# Use a unique project ID for this test run.
# Note: project_id may contain hyphens (used in API calls), but the
# build_keypath helper sanitizes hyphens → underscores for the keypath segment.
RUN_ID = uuid.uuid4().hex[:8]
PROJECT_ID = f"langchain-workflow-{RUN_ID}"

print(f"\n{'='*60}")
print(f"Memstate LangChain Integration - Workflow Test")
print(f"{'='*60}")
print(f"Project ID: {PROJECT_ID}")
print(f"API Base URL: {BASE_URL}")
print(f"{'='*60}\n")


def section(title: str) -> None:
    print(f"\n{'─'*50}")
    print(f"  {title}")
    print(f"{'─'*50}")


def check(label: str, condition: bool, detail: str = "") -> None:
    status = "✓" if condition else "✗"
    print(f"  {status} {label}" + (f": {detail}" if detail else ""))
    if not condition:
        print(f"    FAILED — {label}")


# ---------------------------------------------------------------------------
# Phase 1: API Health Check
# ---------------------------------------------------------------------------

from memstate_langchain import (
    MemstateClient,
    MemstateChatMessageHistory,
    MemstateMemory,
    MemstateSemanticMemory,
    create_memstate_tools,
)

client = MemstateClient(api_key=API_KEY, base_url=BASE_URL)

section("Phase 1: API Health Check")
status = client.get_status()
check("API is reachable", status["status"] == "ok", f"status={status['status']}")
check("LLM processing enabled", status.get("features", {}).get("llm_processing", False))
print(f"  Memory count: {status.get('memory_count', 0)}")
print(f"  Project count: {status.get('project_count', 0)}")


# ---------------------------------------------------------------------------
# Phase 2: Store Initial Project Decisions (direct remember)
# ---------------------------------------------------------------------------

section("Phase 2: Store Initial Project Decisions (memstate_set pattern)")

decisions = [
    ("stack.language.backend", "Python 3.12 with FastAPI", "decision"),
    ("stack.language.frontend", "TypeScript with React 18", "decision"),
    ("stack.database", "PostgreSQL 15 with pgvector", "decision"),
    ("stack.auth", "OAuth2 with Google + JWT RS256", "decision"),
    ("deployment.platform", "Railway for backend, Vercel for frontend", "decision"),
    ("deployment.domain", "api.myapp.com", "fact"),
]

memory_ids = []
for keypath, value, category in decisions:
    # build_keypath automatically sanitizes hyphens in project_id to underscores
    full_keypath = MemstateClient.build_keypath(PROJECT_ID, keypath)
    resp = client.remember(
        content=value,
        keypath=full_keypath,
        project_id=PROJECT_ID,
        category=category,
        source="langchain-test",
    )
    memory_ids.append(resp["memory_id"])
    check(f"Stored {keypath}", "memory_id" in resp, f"id={resp['memory_id'][:8]}...")

print(f"\n  Stored {len(decisions)} decisions.")


# ---------------------------------------------------------------------------
# Phase 3: Ingest Rich Task Summary (server-side intelligence)
# ---------------------------------------------------------------------------

section("Phase 3: Ingest Rich Task Summary (server-side intelligence)")

task_summary = f"""
## Task: Set Up Authentication System

**Project:** {PROJECT_ID}
**Date:** 2025-03-19
**Engineer:** AI Agent

### What Was Done
- Implemented OAuth2 flow with Google as the primary provider
- Added JWT RS256 token signing with 24-hour expiry
- Created middleware for token validation on all protected routes
- Set up refresh token rotation with 30-day sliding window

### Files Changed
- `backend/auth/oauth.py` — Google OAuth2 callback handler
- `backend/auth/jwt.py` — RS256 signing and validation
- `backend/middleware/auth.py` — Request authentication middleware
- `backend/api/users.py` — User profile endpoints

### Key Decisions
- Chose RS256 over HS256 for asymmetric key security
- Refresh tokens stored in Redis with TTL
- Access tokens are stateless (no server-side storage)

### Next Steps
- Add rate limiting to auth endpoints
- Implement MFA with TOTP
- Add audit logging for auth events
"""

ingest_resp = client.ingest(
    project_id=PROJECT_ID,
    content=task_summary,
    source="agent",
    context="authentication system implementation",
)
check("Ingest accepted", "ingestion_id" in ingest_resp or "job_id" in ingest_resp)
check("Status is valid", ingest_resp["status"] in ("complete", "processing", "pending", "duplicate"))
print(f"  Ingestion ID: {ingest_resp.get('ingestion_id', ingest_resp.get('job_id', 'N/A'))}")
print(f"  Status: {ingest_resp['status']}")

if ingest_resp["status"] in ("processing", "pending"):
    print("  Waiting for async ingest to complete...")
    job_id = ingest_resp.get("job_id") or ingest_resp.get("ingestion_id")
    for i in range(15):
        time.sleep(2)
        try:
            job_status = client.get_job_status(job_id)
            if job_status.get("status") in ("complete", "completed", "done"):
                print(f"  ✓ Ingest completed after {(i+1)*2}s")
                break
        except Exception:
            pass
    else:
        print("  ⚠ Ingest still processing (continuing anyway)")
elif ingest_resp["status"] == "complete":
    memories_created = len(ingest_resp.get("memories_created", []))
    keypaths_created = ingest_resp.get("keypaths_created", [])
    print(f"  Memories created: {memories_created}")
    print(f"  Keypaths: {keypaths_created}")

# Wait for search indexing
time.sleep(2)


# ---------------------------------------------------------------------------
# Phase 4: New Session — Retrieve Context Before Starting Work
# ---------------------------------------------------------------------------

section("Phase 4: New Session — Retrieve Context Before Starting Work")

# build_keypath sanitizes hyphens in project_id to underscores for the keypath
project_keypath = MemstateClient.build_keypath(PROJECT_ID)
tree_resp = client.get_keypaths(
    project_id=PROJECT_ID,
    keypath=project_keypath,
    recursive=True,
)
memories = tree_resp.get("memories", {})
check("Project tree is non-empty", len(memories) > 0, f"{len(memories)} memories")
print(f"  Project revision: {tree_resp.get('revision_number', 0)}")
print(f"  Memory tree ({len(memories)} entries):")
for kp, summary in sorted(memories.items())[:8]:
    # Strip the project prefix for display
    short_kp = kp.replace(f"{project_keypath}.", "")
    print(f"    {short_kp}: {str(summary)[:60]}")
if len(memories) > 8:
    print(f"    ... and {len(memories) - 8} more")


# ---------------------------------------------------------------------------
# Phase 5: Semantic Search for Specific Information
# ---------------------------------------------------------------------------

section("Phase 5: Semantic Search for Specific Information")

search_queries = [
    ("authentication token signing", "auth"),
    ("database technology choice", "database"),
    ("deployment platform", "deployment"),
]

for query, expected_domain in search_queries:
    results = client.search(
        query=query,
        project_id=PROJECT_ID,
        limit=3,
    )
    found = results["total_found"] > 0
    check(f"Search '{query}'", found, f"{results['total_found']} results")
    if found:
        top = results["results"][0]
        kp = top.get('keypath', top.get('id', 'unknown'))
        print(f"    Top result: [{kp.split('.')[-1]}] score={top.get('score', 0):.3f}")


# ---------------------------------------------------------------------------
# Phase 6: Update a Decision (Version History)
# ---------------------------------------------------------------------------

section("Phase 6: Update a Decision (Version History)")

# Change the auth decision — this creates a new version
auth_keypath = MemstateClient.build_keypath(PROJECT_ID, "stack.auth")
v1_resp = client.remember(
    content="OAuth2 with Google + JWT RS256",
    keypath=auth_keypath,
    project_id=PROJECT_ID,
    category="decision",
)

v2_resp = client.remember(
    content="OAuth2 with Google + JWT RS256 + MFA with TOTP",
    keypath=auth_keypath,
    project_id=PROJECT_ID,
    category="decision",
)

check("First version stored", v1_resp["version"] >= 1, f"v{v1_resp['version']}")
check("Second version incremented", v2_resp["version"] > v1_resp["version"],
      f"v{v1_resp['version']} → v{v2_resp['version']}")

# Inspect version history
hist = client.history(keypath=auth_keypath, project_id=PROJECT_ID)
check("History has multiple versions", hist["total_versions"] >= 2,
      f"{hist['total_versions']} versions")

latest = next((v for v in hist["versions"] if v.get("is_latest")), None)
check("Latest version is marked", latest is not None)
if latest:
    print(f"  Latest: {latest.get('summary', latest.get('content', ''))[:80]}")


# ---------------------------------------------------------------------------
# Phase 7: MemstateSemanticMemory (modern LCEL pattern)
# ---------------------------------------------------------------------------

section("Phase 7: MemstateSemanticMemory (modern LCEL pattern)")

memory = MemstateSemanticMemory(
    api_key=API_KEY,
    project_id=PROJECT_ID,
    session_id="workflow-test",
    base_url=BASE_URL,
    search_limit=5,
)

# Test retrieve()
enriched = memory.retrieve({"input": "What authentication system are we using?"})
check("retrieve() returns dict with memory_key", "relevant_memories" in enriched)
check("retrieve() preserves input key", "input" in enriched)
has_content = bool(enriched.get("relevant_memories"))
print(f"  Memory context: {'present (' + str(len(enriched['relevant_memories'])) + ' chars)' if has_content else 'empty (search may still be indexing)'}")
if has_content:
    print(f"  Preview: {enriched['relevant_memories'][:200]!r}")

# Test save()
memory.save(
    inputs={"input": "What authentication system are we using?"},
    output="We are using OAuth2 with Google + JWT RS256 + MFA with TOTP.",
)
print("  ✓ save() completed without error")

# Test wrap() with a mock chain
from langchain_core.runnables import RunnableLambda

def mock_llm(inputs: dict) -> str:
    user_input = inputs.get("input", "")
    memories = inputs.get("relevant_memories", "")
    if "database" in user_input.lower():
        return "PostgreSQL 15 with pgvector."
    if "auth" in user_input.lower():
        return "OAuth2 with Google + JWT RS256."
    return f"Context available: {len(memories)} chars."

mock_chain = RunnableLambda(mock_llm)
memory2 = MemstateSemanticMemory(
    api_key=API_KEY,
    project_id=PROJECT_ID,
    session_id="workflow-wrap-test",
    base_url=BASE_URL,
    search_limit=5,
)
chain_with_memory = memory2.wrap(mock_chain)

captured_inputs = {}
original_retrieve = memory2.retrieve
def capturing_retrieve(inputs):
    result = original_retrieve(inputs)
    captured_inputs.update(result)
    return result
memory2.retrieve = capturing_retrieve

response = chain_with_memory.invoke({"input": "What database are we using?"})
check("wrap() returns chain output", isinstance(response, str))
print(f"  Chain response: {response!r}")

# Test get_memories()
search_results = memory.get_memories(query="authentication", limit=5)
check("get_memories(query=...) returns list", isinstance(search_results, list))
print(f"  get_memories(query='authentication'): {len(search_results)} results")

all_memories = memory.get_memories(limit=20)
check("get_memories() (browse) returns list", isinstance(all_memories, list))
print(f"  get_memories() (browse): {len(all_memories)} memories")


# ---------------------------------------------------------------------------
# Phase 8: MemstateChatMessageHistory
# ---------------------------------------------------------------------------

section("Phase 8: MemstateChatMessageHistory")

from langchain_core.messages import AIMessage, HumanMessage

session_id = f"workflow-chat-{RUN_ID}"
chat_history = MemstateChatMessageHistory(
    api_key=API_KEY,
    project_id=PROJECT_ID,
    session_id=session_id,
    base_url=BASE_URL,
)

# Start with empty history
initial_messages = chat_history.messages
check("Initial history is empty", len(initial_messages) == 0,
      f"{len(initial_messages)} messages")

# Add a conversation
chat_history.add_message(HumanMessage(content="What stack are we using?"))
chat_history.add_message(AIMessage(content="We're using FastAPI + PostgreSQL + React."))
chat_history.add_message(HumanMessage(content="What about auth?"))
chat_history.add_message(AIMessage(content="OAuth2 with Google and JWT RS256 tokens."))

# Retrieve
messages = chat_history.messages
check("Messages persisted", len(messages) == 4, f"{len(messages)} messages")
check("First message is HumanMessage", isinstance(messages[0], HumanMessage))
check("Second message is AIMessage", isinstance(messages[1], AIMessage))
check("Content preserved", "FastAPI" in messages[1].content)

# Simulate a new session (new object, same session_id)
chat_history2 = MemstateChatMessageHistory(
    api_key=API_KEY,
    project_id=PROJECT_ID,
    session_id=session_id,
    base_url=BASE_URL,
)
reloaded = chat_history2.messages
check("Messages reload in new session", len(reloaded) == 4, f"{len(reloaded)} messages")

# Test add_messages (bulk)
chat_history.add_messages([
    HumanMessage(content="What's the deployment setup?"),
    AIMessage(content="Railway for backend, Vercel for frontend."),
])
messages_after_bulk = chat_history.messages
check("add_messages() appends correctly", len(messages_after_bulk) == 6,
      f"{len(messages_after_bulk)} messages")

# Test clear
chat_history.clear()
cleared_messages = chat_history.messages
check("clear() removes messages", len(cleared_messages) == 0,
      f"{len(cleared_messages)} messages remaining")


# ---------------------------------------------------------------------------
# Phase 9: Agent Tools
# ---------------------------------------------------------------------------

section("Phase 9: Agent Tools")

tools = create_memstate_tools(
    api_key=API_KEY,
    project_id=PROJECT_ID,
    base_url=BASE_URL,
)

check("All 6 tools created", len(tools) == 6, f"{len(tools)} tools")
tool_names = {t.name for t in tools}
for expected in ["memstate_remember", "memstate_set", "memstate_get",
                 "memstate_search", "memstate_history", "memstate_delete"]:
    check(f"Tool '{expected}' present", expected in tool_names)

# Test each tool
get_tool = next(t for t in tools if t.name == "memstate_get")
get_result = get_tool._run()
check("memstate_get returns content", "Error" not in get_result, get_result[:60])

search_tool = next(t for t in tools if t.name == "memstate_search")
search_result = search_tool._run(query="authentication")
check("memstate_search returns content", "Error" not in search_result, search_result[:60])

set_tool = next(t for t in tools if t.name == "memstate_set")
set_result = set_tool._run(keypath="test.tool.value", value="tool_test_123")
check("memstate_set returns content", "Error" not in set_result, set_result[:60])

history_tool = next(t for t in tools if t.name == "memstate_history")
history_result = history_tool._run(keypath="stack.auth")
check("memstate_history returns content", "Error" not in history_result, history_result[:60])

delete_tool = next(t for t in tools if t.name == "memstate_delete")
delete_result = delete_tool._run(keypath="test.tool.value")
check("memstate_delete returns content", "Error" not in delete_result, delete_result[:60])

remember_tool = next(t for t in tools if t.name == "memstate_remember")
remember_result = remember_tool._run(
    content="## Tool Test\nAll tools are working correctly.",
    source="langchain-test",
)
check("memstate_remember returns content", "Error" not in remember_result, remember_result[:60])


# ---------------------------------------------------------------------------
# Phase 10: Optional LangChain Agent Test (requires OpenAI)
# ---------------------------------------------------------------------------

if OPENAI_API_KEY:
    section("Phase 10: Full LangChain Agent with Memstate Tools (OpenAI)")

    try:
        from langchain_openai import ChatOpenAI
        from langchain.agents import create_tool_calling_agent, AgentExecutor
        from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder

        llm = ChatOpenAI(model="gpt-4o-mini", temperature=0, api_key=OPENAI_API_KEY)

        prompt = ChatPromptTemplate.from_messages([
            ("system", f"""You are a helpful assistant with persistent memory powered by Memstate AI.

Before answering questions, use memstate_get to check what you already know about the project.
After completing tasks, use memstate_remember to save what you learned.

Current project: {PROJECT_ID}
"""),
            MessagesPlaceholder(variable_name="chat_history"),
            ("human", "{input}"),
            MessagesPlaceholder(variable_name="agent_scratchpad"),
        ])

        agent = create_tool_calling_agent(llm, tools, prompt)
        executor = AgentExecutor(
            agent=agent,
            tools=tools,
            verbose=False,
            max_iterations=5,
            handle_parsing_errors=True,
        )

        # Test 1: Ask about the project (should use memstate_get)
        response1 = executor.invoke({
            "input": "What database are we using in this project?",
            "chat_history": [],
        })
        check("Agent answered database question",
              bool(response1.get("output")),
              response1.get("output", "")[:60])

        # Test 2: Ask the agent to remember a new decision
        response2 = executor.invoke({
            "input": "Remember that we decided to add Redis for caching with a 1-hour TTL.",
            "chat_history": [],
        })
        check("Agent saved new decision",
              bool(response2.get("output")),
              response2.get("output", "")[:60])

        print("\n  ✓ Full LangChain agent test passed!")

    except ImportError as e:
        print(f"  ⚠ Skipping agent test: {e}")
    except Exception as e:
        print(f"  ✗ Agent test failed: {e}")
else:
    section("Phase 10: Full LangChain Agent Test (SKIPPED — no OPENAI_API_KEY)")
    print("  Set OPENAI_API_KEY to run the full agent test.")


# ---------------------------------------------------------------------------
# Cleanup
# ---------------------------------------------------------------------------

section("Cleanup")

try:
    client.delete_memory(
        keypath=MemstateClient.build_keypath(PROJECT_ID),
        project_id=PROJECT_ID,
        recursive=True,
    )
    print(f"  ✓ Cleaned up project '{PROJECT_ID}'")
except Exception as e:
    print(f"  ⚠ Cleanup failed (non-critical): {e}")


# ---------------------------------------------------------------------------
# Summary
# ---------------------------------------------------------------------------

print(f"\n{'='*60}")
print("  Workflow Test Complete!")
print(f"{'='*60}\n")
