from agentgit.database.repositories.user_repository import UserRepository
from agentgit.auth.user import User
from datetime import datetime

user_repo = UserRepository()  # Auto-creates rootusr

# Option A: Use default user
user = user_repo.find_by_username("rootusr")

# Option B: Create custom user
new_user = User(username="alice", created_at=datetime.now())
new_user.set_password("secure_password")
user = user_repo.save(new_user)

from agentgit.agents.rollback_agent import RollbackAgent
from agentgit.sessions.external_session import ExternalSession
from agentgit.database.repositories.external_session_repository import ExternalSessionRepository
from agentgit.database.repositories.user_repository import UserRepository
from langchain_openai import ChatOpenAI
from langchain_core.tools import tool
from datetime import datetime

# 1. Get user
user_repo = UserRepository()
user = user_repo.find_by_username("rootusr")

# 2. Create session
external_repo = ExternalSessionRepository()
session = external_repo.create(ExternalSession(
    user_id=user.id,
    session_name="My First Agent",
    created_at=datetime.now()
))

# 3. Define tools
@tool
def calculate_sum(a: int, b: int) -> int:
    """Add two numbers together."""
    return a + b

# 4. Create agent
agent = RollbackAgent(
    external_session_id=session.id,
    model=ChatOpenAI(model="gpt-4o-mini"),
    tools=[calculate_sum],
    auto_checkpoint=True
)

# 5. Chat!
response = agent.run("What is 25 + 17?")
print(response)

from agentgit.agents.rollback_agent import RollbackAgent
from agentgit.sessions.external_session import ExternalSession
from agentgit.database.repositories.external_session_repository import ExternalSessionRepository
from agentgit.database.repositories.checkpoint_repository import CheckpointRepository
from agentgit.database.repositories.internal_session_repository import InternalSessionRepository
from agentgit.database.repositories.user_repository import UserRepository
from agentgit.auth.user import User
from langchain_openai import ChatOpenAI
from langchain_core.tools import tool
from datetime import datetime

# Step 1: Create or get a user
user_repo = UserRepository()  # Auto-creates 'rootusr' with ID=1
user = user_repo.find_by_username("rootusr")

# Step 2: Setup repositories
external_repo = ExternalSessionRepository()
checkpoint_repo = CheckpointRepository()
internal_repo = InternalSessionRepository()

external_session = external_repo.create(ExternalSession(
    user_id=user.id,
    session_name="Rollback Demo",
    created_at=datetime.now()
))

# Define a tool with side effects
order_database = {}  # Simulated database

@tool
def create_order(order_id: str, amount: float) -> str:
    """Create a new order in the database."""
    order_database[order_id] = {"amount": amount, "status": "created"}
    return f"Order {order_id} created for ${amount}"

def reverse_create_order(args, result):
    """Reverse order creation by deleting it."""
    order_id = args['order_id']
    if order_id in order_database:
        del order_database[order_id]
    return f"Order {order_id} deleted"

# Create agent with rollback capability
agent = RollbackAgent(
    external_session_id=external_session.id,
    model=ChatOpenAI(model="gpt-4o-mini"),
    tools=[create_order],
    reverse_tools={"create_order": reverse_create_order},
    auto_checkpoint=True,
    checkpoint_repo=checkpoint_repo,
    internal_session_repo=internal_repo
)

# Conversation flow
print("=== Original Timeline ===")
agent.run("Hello! I need to create some orders.")

# Create manual checkpoint before operations
checkpoint_msg = agent.create_checkpoint_tool("Before order creation")
print(checkpoint_msg)
# Extract checkpoint ID from message: "✓ Checkpoint 'Before order creation' created successfully (ID: 1)"
safe_checkpoint_id = int(checkpoint_msg.split("ID: ")[1].rstrip(")")) # In practice, parse this from checkpoint_msg

# Create some orders (auto-checkpoints created after each)
agent.run("Please create order #1001 for $250")
print(f"Database state: {order_database}")

agent.run("Now create order #1002 for $150")
print(f"Database state: {order_database}")

agent.run("Create one more order #1003 for $300")
print(f"Database state: {order_database}")
# Database now has: {1001: ..., 1002: ..., 1003: ...}

# Get conversation history length
original_history = agent.get_conversation_history()
print(f"Original history length: {len(original_history)} messages")

# Oops! We need to go back to before the orders were created
print("\n=== Performing Rollback ===")

# Rollback to the safe checkpoint - creates a new branch
branched_agent = RollbackAgent.from_checkpoint(
    checkpoint_id=safe_checkpoint_id,
    external_session_id=external_session.id,
    model=ChatOpenAI(model="gpt-4o-mini"),
    tools=[create_order],
    reverse_tools={"create_order": reverse_create_order},
    checkpoint_repo=checkpoint_repo,
    internal_session_repo=internal_repo
)

# Check the branched agent's state
branched_history = branched_agent.get_conversation_history()
print(f"Branched history length: {len(branched_history)} messages")
print(f"Database state after branch: {order_database}")

# The branch has the state from before the checkpoint
# But the tool operations haven't been reversed yet
# To reverse tools, get the checkpoint and rollback from its track position
checkpoint = checkpoint_repo.get_by_id(safe_checkpoint_id)
if "tool_track_position" in checkpoint.metadata:
    track_position = checkpoint.metadata["tool_track_position"]
    reverse_results = agent.rollback_tools_from_track_index(track_position)

    print("\nTool Reversal Results:")
    for result in reverse_results:
        if result.reversed_successfully:
            print(f"✓ Reversed {result.tool_name}")
        else:
            print(f"✗ Failed to reverse {result.tool_name}: {result.error_message}")

print(f"Database after tool reversal: {order_database}")

# Continue with different actions on the branch
print("\n=== New Branch Timeline ===")
branched_agent.run("Let's create just one order instead: #2001 for $500")
print(f"Database state: {order_database}")

# Original agent can still continue independently
print("\n=== Original Timeline Continues ===")
agent.run("Create another order #1004 for $200")

# Show conversation histories to prove independence
print("\n=== Timeline Independence Verification ===")
original_conv = agent.get_conversation_history()
branched_conv = branched_agent.get_conversation_history()

print(f"\nOriginal timeline - conversation history ({len(original_conv)} messages):")
for i, msg in enumerate(original_conv[-4:], start=len(original_conv)-3):  # Show last 4 messages
    role = msg.get('role', 'unknown')
    content = msg.get('content', '')[:80]  # Truncate long messages
    print(f"  [{i}] {role}: {content}...")

print(f"\nBranched timeline - conversation history ({len(branched_conv)} messages):")
for i, msg in enumerate(branched_conv, start=1):
    role = msg.get('role', 'unknown')
    content = msg.get('content', '')[:80]
    print(f"  [{i}] {role}: {content}...")

# Verify both sessions exist
all_sessions = internal_repo.get_by_external_session(external_session.id)
print(f"\nTotal internal sessions (branches): {len(all_sessions)}")
print(f"Original session ID: {agent.internal_session.id}")
print(f"Branched session ID: {branched_agent.internal_session.id}")
print(f"Branch parent: {branched_agent.internal_session.parent_session_id}")
print(f"Is branch: {branched_agent.internal_session.is_branch()}")

from agentgit.agents.agent_service import AgentService
from agentgit.sessions.external_session import ExternalSession
from agentgit.database.repositories.external_session_repository import ExternalSessionRepository
from agentgit.database.repositories.user_repository import UserRepository
from langchain_core.tools import tool
from datetime import datetime

# Define tools
@tool
def process_payment(amount: float, order_id: str) -> str:
    """Process a payment for an order."""
    return f"Payment of ${amount} processed for order {order_id}"

def reverse_payment(args, result):
    """Reverse a payment."""
    print(f"  ✓ Reversed: Refunded ${args['amount']} for order {args['order_id']}")
    return f"Refunded ${args['amount']} for order {args['order_id']}"

# Step 1: Create user and external session (same as before)
user_repo = UserRepository()
user = user_repo.find_by_username("rootusr")

external_repo = ExternalSessionRepository()
external_session = external_repo.create(ExternalSession(
    user_id=user.id,
    session_name="Payment Processing",
    created_at=datetime.now()
))

# Step 2: Initialize service (auto-creates all repositories)
service = AgentService()
# ✓ No manual repository creation
# ✓ Model config loaded from environment
# ✓ All dependencies managed internally

# Step 3: Create agent with minimal code
agent = service.create_new_agent(
    external_session_id=external_session.id,
    tools=[process_payment],
    reverse_tools={"process_payment": reverse_payment},
    auto_checkpoint=True
)
# ✓ No model creation needed
# ✓ No repository passing needed
# ✓ Best-practice defaults applied

# Step 4: Use the agent
response1 = agent.run("Process a payment of $150 for order #1001")
print(response1)

# Create checkpoint
checkpoint_msg = agent.create_checkpoint_tool("Before second payment")
checkpoint_id = int(checkpoint_msg.split("ID: ")[1].split(")")[0])

response2 = agent.run("Process another payment of $250 for order #1002")
print(response2)

# Step 5: List checkpoints using service utility
checkpoints = service.list_checkpoints(agent.internal_session.id)
print(f"Total checkpoints: {len(checkpoints)}")

# Step 6: Resume session later
# Simulate app restart
service = AgentService()
resumed_agent = service.resume_agent(
    external_session_id=external_session.id,
    tools=[process_payment],
    reverse_tools={"process_payment": reverse_payment}
)
# ✓ Conversation history automatically restored
# ✓ Session state automatically loaded
# ✓ Ready to continue immediately

if resumed_agent:
    response3 = resumed_agent.run("What payments did we process?")
    print(response3)

# Step 7: Rollback with automatic tool reversal
rolled_back = service.rollback_to_checkpoint(
    external_session_id=external_session.id,
    checkpoint_id=checkpoint_id,
    rollback_tools=True,  # Automatically reverses payment operations!
    tools=[process_payment],
    reverse_tools={"process_payment": reverse_payment}
)
# ✓ Checkpoint retrieved automatically
# ✓ Tool operations reversed automatically
# ✓ New branch created automatically
# ✓ Ready to use immediately

if rolled_back:
    response4 = rolled_back.run("Let's try a different payment amount: $300 for order #1003")
    print(response4)
    print()

# Automatic checkpoints (after tool calls)
agent = RollbackAgent(..., auto_checkpoint=True)

# Manual checkpoints
checkpoint_msg = agent.create_checkpoint_tool("Before risky operation")

# List all checkpoints
agent.list_checkpoints_tool()

# Get checkpoint details
agent.get_checkpoint_info_tool(checkpoint_id=1)

# Rollback creates a new branch (preserves original timeline)
new_agent = RollbackAgent.from_checkpoint(
    checkpoint_id=checkpoint_id,
    external_session_id=session.id,
    model=model,
    tools=tools,
    reverse_tools=reverse_tools,
    checkpoint_repo=checkpoint_repo,
    internal_session_repo=internal_repo
)

# Or use AgentService for one-liner
rolled_back = service.rollback_to_checkpoint(
    external_session_id=session.id,
    checkpoint_id=checkpoint_id,
    rollback_tools=True
)

# Define tool
@tool
def save_record(record_id: str, data: str) -> str:
    """Save data to database."""
    db.write(record_id, data)
    return f"Saved {record_id}"

# Define reverse function
def reverse_save(args, result):
    """Delete the saved record."""
    db.delete(args['record_id'])
    return f"Deleted {args['record_id']}"

# Register when creating agent
agent = RollbackAgent(
    ...,
    tools=[save_record],
    reverse_tools={"save_record": reverse_save}
)

# Check if session is a branch
agent.internal_session.is_branch()

# Get parent session ID
agent.internal_session.parent_session_id

# Get branch point checkpoint
agent.internal_session.branch_point_checkpoint_id

# List all branches for a session
branches = internal_repo.get_by_external_session(session.id)
print(f"Total branches: {len(branches)}")

# Agent helps customer, makes mistake
agent.run("Process refund of $500")

# Realize error, rollback to before refund
checkpoint_id = 5  # From before refund
corrected = service.rollback_to_checkpoint(
    external_session_id=session.id,
    checkpoint_id=checkpoint_id,
    rollback_tools=True  # Reverses the refund!
)

# Try again with correct amount
corrected.run("Process refund of $50")

# Create checkpoint after gathering context
checkpoint_msg = agent.create_checkpoint_tool("Context gathered")
checkpoint_id = 1

# Test approach A
branch_a = RollbackAgent.from_checkpoint(checkpoint_id, ...)
response_a = branch_a.run("Explain technically")

# Test approach B
branch_b = RollbackAgent.from_checkpoint(checkpoint_id, ...)
response_b = branch_b.run("Explain simply")

# Compare results
print(f"Technical: {len(response_a)} chars")
print(f"Simple: {len(response_b)} chars")

# Create checkpoint before risky operation
agent.create_checkpoint_tool("Before deletion")

# Try risky operation
agent.run("Delete all test records")

# If something goes wrong, rollback
service.rollback_to_checkpoint(
    external_session_id=session.id,
    checkpoint_id=checkpoint_id,
    rollback_tools=True  # Undoes deletions!
)

# Conversation
agent.run(message: str) -> str

# Checkpoints
agent.create_checkpoint_tool(name: str) -> str
agent.list_checkpoints_tool() -> str
agent.get_checkpoint_info_tool(checkpoint_id: int) -> str
agent.cleanup_auto_checkpoints_tool(keep_latest: int) -> str

# State inspection
agent.get_conversation_history() -> List[dict]
agent.get_session_state() -> dict
agent.get_tool_track() -> List[ToolInvocationRecord]

# Rollback
RollbackAgent.from_checkpoint(checkpoint_id, ...) -> RollbackAgent
agent.rollback_tools_from_track_index(position: int) -> List[ReverseResult]

# Agent management
service.create_new_agent(external_session_id, tools, ...) -> RollbackAgent
service.resume_agent(external_session_id) -> RollbackAgent | None

# Rollback
service.rollback_to_checkpoint(
    external_session_id,
    checkpoint_id,
    rollback_tools=True
) -> RollbackAgent | None

# Utilities
service.list_checkpoints(internal_session_id) -> List[Checkpoint]