"""
Workspace Artifacts Example

This example demonstrates using Workspace for artifact sharing
and state management across agents.

Setup:
1. Install: pip install synqed openai python-dotenv
2. Create .env file with: OPENAI_API_KEY='your-key-here'
3. Run: python workspace_artifacts.py
"""

import asyncio
import os
from pathlib import Path

from dotenv import load_dotenv

import synqed

# Load environment
load_dotenv()
load_dotenv(dotenv_path=Path(__file__).parent / ".env")


async def main():
    # Check for API key
    api_key = os.getenv("OPENAI_API_KEY")
    if not api_key:
        print("❌ Please set OPENAI_API_KEY in your .env file")
        return

    print("\n" + "=" * 70)
    print("  WORKSPACE ARTIFACTS & STATE EXAMPLE")
    print("=" * 70 + "\n")

    # Step 1: Create agents
    print("📦 Creating agents...\n")

    async def data_agent_executor(context):
        """Data processing agent."""
        user_message = context.get_user_input()
        from openai import AsyncOpenAI

        client = AsyncOpenAI(api_key=api_key)
        response = await client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[
                {
                    "role": "system",
                    "content": "You are a data analysis expert. Process and analyze data, providing insights.",
                },
                {"role": "user", "content": user_message},
            ],
        )
        return response.choices[0].message.content

    async def report_agent_executor(context):
        """Report generation agent."""
        user_message = context.get_user_input()
        from openai import AsyncOpenAI

        client = AsyncOpenAI(api_key=api_key)
        response = await client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[
                {
                    "role": "system",
                    "content": "You are a report writer. Create professional reports from data and analysis.",
                },
                {"role": "user", "content": user_message},
            ],
        )
        return response.choices[0].message.content

    data_agent = synqed.Agent(
        name="DataAgent",
        description="Processes and analyzes data",
        skills=["data_processing", "analysis"],
        executor=data_agent_executor,
    )

    report_agent = synqed.Agent(
        name="ReportAgent",
        description="Creates professional reports",
        skills=["report_generation", "documentation"],
        executor=report_agent_executor,
    )

    print("✓ Created DataAgent")
    print("✓ Created ReportAgent\n")

    # Step 2: Start servers
    print("🚀 Starting servers...\n")

    data_server = synqed.AgentServer(data_agent, port=8001)
    report_server = synqed.AgentServer(report_agent, port=8002)

    await data_server.start_background()
    await report_server.start_background()

    print(f"✓ Servers running\n")

    # Step 3: Create workspace with persistence
    print("🏗️  Creating workspace with persistence...\n")

    workspace = synqed.Workspace(
        name="Data Analysis Workspace",
        description="Workspace for data processing and report generation",
        auto_cleanup=False,  # Keep artifacts
        enable_persistence=True,  # Save state
    )

    workspace.add_agent(data_agent)
    workspace.add_agent(report_agent)

    await workspace.start()
    print(f"✓ Workspace: {workspace.name}\n")

    # Step 4: Set shared state
    print("=" * 70)
    print("SHARED STATE MANAGEMENT")
    print("=" * 70 + "\n")

    # Set some shared state that agents can access
    workspace.set_shared_state("project_name", "Q4 Sales Analysis")
    workspace.set_shared_state("deadline", "2025-12-31")
    workspace.set_shared_state("department", "Sales")

    print("✓ Set shared state:")
    for key, value in workspace.get_all_shared_state().items():
        print(f"  • {key}: {value}")
    print()

    # Step 5: Create artifacts
    print("=" * 70)
    print("ARTIFACT CREATION")
    print("=" * 70 + "\n")

    # Sample data artifact
    sales_data = {
        "q1": {"revenue": 120000, "units": 450},
        "q2": {"revenue": 135000, "units": 480},
        "q3": {"revenue": 142000, "units": 510},
        "q4": {"revenue": 158000, "units": 540},
    }

    artifact_id = workspace.add_artifact(
        name="sales_data.json",
        artifact_type="data",
        content=sales_data,
        created_by="system",
        metadata={"source": "sales_database", "year": 2025},
    )

    print(f"✓ Created data artifact: {artifact_id}\n")

    # Create a markdown file artifact
    readme_content = """# Sales Analysis Project

## Overview
This workspace contains the Q4 sales analysis.

## Data Sources
- Internal sales database
- Customer feedback forms

## Team
- DataAgent: Processes raw data
- ReportAgent: Generates final reports
"""

    workspace.add_artifact(
        name="README.md",
        artifact_type="file",
        content=readme_content,
        created_by="system",
        metadata={"version": "1.0"},
    )

    print("✓ Created README artifact\n")

    # Step 6: Agents collaborate with access to artifacts
    print("=" * 70)
    print("COLLABORATIVE ANALYSIS")
    print("=" * 70 + "\n")

    # Get project info from shared state
    project = workspace.get_shared_state("project_name")
    print(f"Project: {project}\n")

    task = f"""
    Analyze the sales data and provide insights.
    Available data: {sales_data}
    Project: {project}
    """

    print("🤖 Agents processing...\n")
    results = await workspace.collaborate(task.strip())

    for agent_name, response in results.items():
        print(f"📝 {agent_name}:")
        print(f"{response}\n")
        print("-" * 70 + "\n")

    # Step 7: View all artifacts
    print("=" * 70)
    print("WORKSPACE ARTIFACTS")
    print("=" * 70 + "\n")

    artifacts = workspace.get_artifacts()
    print(f"Total artifacts: {len(artifacts)}\n")

    for artifact in artifacts:
        print(f"📦 {artifact.name}")
        print(f"   Type: {artifact.artifact_type}")
        print(f"   Created: {artifact.created_at.strftime('%Y-%m-%d %H:%M:%S')}")
        print(f"   By: {artifact.created_by}")
        if artifact.metadata:
            print(f"   Metadata: {artifact.metadata}")
        print()

    # Step 8: Filter artifacts by type
    print("=" * 70)
    print("ARTIFACT FILTERING")
    print("=" * 70 + "\n")

    data_artifacts = workspace.get_artifacts(artifact_type="data")
    result_artifacts = workspace.get_artifacts(artifact_type="result")

    print(f"Data artifacts: {len(data_artifacts)}")
    print(f"Result artifacts: {len(result_artifacts)}\n")

    # Step 9: Broadcast message to all agents
    print("=" * 70)
    print("BROADCAST MESSAGE")
    print("=" * 70 + "\n")

    broadcast_msg = "Please summarize your contributions to this project"
    print(f"Broadcasting: {broadcast_msg}\n")

    responses = await workspace.broadcast_message(broadcast_msg)

    for agent_name, response in responses.items():
        print(f"📝 {agent_name}:")
        print(f"   {response[:150]}...")
        print()

    # Step 10: Export workspace
    print("=" * 70)
    print("WORKSPACE EXPORT")
    print("=" * 70 + "\n")

    export_path = await workspace.export_workspace()
    print(f"✓ Workspace exported to:\n  {export_path}\n")

    print(f"📂 Workspace directory:\n  {workspace.workspace_dir}\n")

    # Step 11: Clean up
    await workspace.complete()
    await workspace.close()

    await data_server.stop()
    await report_server.stop()

    print("=" * 70)
    print("✓ Example completed successfully!")
    print(f"✓ Artifacts saved in: {workspace.workspace_dir}")
    print("=" * 70 + "\n")


if __name__ == "__main__":
    asyncio.run(main())

