#!/usr/bin/env python3
"""
Temporal Workflows API Client Examples
Demonstrates how to use the /workflows endpoint with Python
"""

import requests
import json
import time
from typing import Dict, Any, List, Optional

GATEWAY = "https://api.riotpiao.com"


class WorkflowClient:
    """Simple client for interacting with the Workflows API"""

    def __init__(self, base_url: str = GATEWAY):
        self.base_url = base_url
        self.session = requests.Session()

    def execute_workflow(
        self,
        workflow: str,
        input_data: Dict[str, Any],
        timeout: Optional[int] = None,
        wait: bool = True,
    ) -> Dict[str, Any]:
        """
        Execute a workflow

        Args:
            workflow: Workflow name
            input_data: Input parameters for the workflow
            timeout: Timeout in seconds (default: 30)
            wait: Whether to wait for completion (default: True)

        Returns:
            Workflow response dict with status, output, etc.
        """
        payload = {
            "workflow": workflow,
            "input": input_data,
        }

        if timeout is not None:
            payload["timeout"] = timeout

        if not wait:
            payload["wait"] = False

        response = self.session.post(
            f"{self.base_url}/workflows",
            json=payload,
            headers={"Content-Type": "application/json"},
        )

        response.raise_for_status()
        return response.json()

    def chat_and_embed(
        self, model: str, messages: List[Dict[str, str]], embed_model: Optional[str] = None
    ) -> Dict[str, Any]:
        """
        Chat with a model and embed the response

        Args:
            model: Chat model name
            messages: Messages in OpenAI format
            embed_model: Optional embedding model (default: nomic-ai/nomic-embed-text-v2-moe)

        Returns:
            Workflow response with chat and embedding results
        """
        input_data = {
            "model": model,
            "messages": messages,
        }

        if embed_model:
            input_data["embed_model"] = embed_model

        return self.execute_workflow("chat-and-embed", input_data)

    def multi_model_chat(self, models: List[str], messages: List[Dict[str, str]]) -> Dict[str, Any]:
        """
        Chat with multiple models and compare responses

        Args:
            models: List of model names
            messages: Messages in OpenAI format

        Returns:
            Workflow response with results from all models
        """
        return self.execute_workflow(
            "multi-model-chat",
            {
                "models": models,
                "messages": messages,
            },
        )

    def rag_pipeline(
        self,
        query: str,
        documents: List[str],
        model: Optional[str] = None,
        rerank_model: Optional[str] = None,
        top_k: int = 3,
    ) -> Dict[str, Any]:
        """
        RAG pipeline: rerank documents and answer based on top results

        Args:
            query: User query or question
            documents: List of document texts
            model: Chat model (default: "reasoning")
            rerank_model: Reranker model (default: "BAAI/bge-reranker-base")
            top_k: Number of top documents to use (default: 3)

        Returns:
            Workflow response with reranked documents and chat answer
        """
        input_data = {
            "query": query,
            "documents": documents,
            "top_k": top_k,
        }

        if model:
            input_data["model"] = model

        if rerank_model:
            input_data["rerank_model"] = rerank_model

        return self.execute_workflow("rag-pipeline", input_data)

    def batch_embeddings(
        self, texts: List[str], model: Optional[str] = None
    ) -> Dict[str, Any]:
        """
        Generate embeddings for multiple texts

        Args:
            texts: List of text strings
            model: Embedding model (default: nomic-ai/nomic-embed-text-v2-moe)

        Returns:
            Workflow response with embedding results
        """
        input_data = {"texts": texts}

        if model:
            input_data["model"] = model

        return self.execute_workflow("batch-embeddings", input_data)


def example_chat_and_embed():
    """Example: Chat and embed"""
    print("\n" + "="*50)
    print("Example 1: Chat and Embed")
    print("="*50)

    client = WorkflowClient()
    result = client.chat_and_embed(
        model="reasoning",
        messages=[
            {
                "role": "user",
                "content": "What is machine learning in one sentence?",
            }
        ],
    )

    print(f"Workflow ID: {result['id']}")
    print(f"Status: {result['status']}")
    print(f"Chat Response: {result['output']['chat_response']['choices'][0]['message']['content']}")
    print(f"Embedding dimensions: {len(result['output']['embedding_response']['data'][0]['embedding'])}")


def example_multi_model_chat():
    """Example: Multi-model chat"""
    print("\n" + "="*50)
    print("Example 2: Multi-Model Chat")
    print("="*50)

    client = WorkflowClient()
    result = client.multi_model_chat(
        models=["reasoning", "ornith:35b"],
        messages=[
            {
                "role": "user",
                "content": "What is the capital of France?",
            }
        ],
    )

    print(f"Workflow ID: {result['id']}")
    print(f"Status: {result['status']}")

    for model_result in result["output"]:
        model = model_result["model"]
        if "result" in model_result:
            content = model_result["result"]["choices"][0]["message"]["content"]
            print(f"\n{model}: {content}")
        elif "error" in model_result:
            print(f"\n{model}: Error - {model_result['error']}")


def example_rag_pipeline():
    """Example: RAG pipeline"""
    print("\n" + "="*50)
    print("Example 3: RAG Pipeline")
    print("="*50)

    client = WorkflowClient()
    result = client.rag_pipeline(
        query="How does photosynthesis work?",
        documents=[
            "Photosynthesis is the process by which plants convert sunlight into chemical energy.",
            "The mitochondria is the powerhouse of the cell.",
            "Light reactions occur in the thylakoid membrane of chloroplasts.",
            "Dogs are domesticated animals.",
            "The Calvin cycle produces glucose from CO2.",
        ],
        top_k=2,
    )

    print(f"Workflow ID: {result['id']}")
    print(f"Status: {result['status']}")
    print(f"\nTop Documents:")
    for i, doc in enumerate(result["output"]["reranked_documents"], 1):
        print(f"  {i}. {doc[:80]}...")

    print(f"\nChat Response:")
    print(f"  {result['output']['chat_response']['choices'][0]['message']['content'][:200]}...")


def example_batch_embeddings():
    """Example: Batch embeddings"""
    print("\n" + "="*50)
    print("Example 4: Batch Embeddings")
    print("="*50)

    client = WorkflowClient()
    result = client.batch_embeddings(
        texts=[
            "The quick brown fox",
            "Machine learning is powerful",
            "Python is a great language",
        ]
    )

    print(f"Workflow ID: {result['id']}")
    print(f"Status: {result['status']}")
    print(f"Number of embeddings: {len(result['output']['data'])}")
    print(f"Embedding dimensions: {len(result['output']['data'][0]['embedding'])}")
    print(f"Model used: {result['output']['model']}")


def example_error_handling():
    """Example: Error handling"""
    print("\n" + "="*50)
    print("Example 5: Error Handling")
    print("="*50)

    client = WorkflowClient()

    # Try unknown workflow
    print("\nAttempting unknown workflow...")
    try:
        result = client.execute_workflow("nonexistent", {})
        if result.get("status") == "failed":
            print(f"Workflow failed: {result.get('error')}")
        else:
            print(f"Response: {json.dumps(result, indent=2)}")
    except requests.exceptions.HTTPError as e:
        print(f"HTTP Error: {e}")
        print(f"Response: {e.response.json()}")

    # Try missing required parameter
    print("\nAttempting chat-and-embed without model...")
    try:
        result = client.execute_workflow("chat-and-embed", {"messages": []})
        if result.get("status") == "failed":
            print(f"Workflow failed: {result.get('error')}")
    except requests.exceptions.HTTPError as e:
        print(f"HTTP Error: {e}")


def example_custom_timeout():
    """Example: Custom timeout"""
    print("\n" + "="*50)
    print("Example 6: Custom Timeout")
    print("="*50)

    client = WorkflowClient()
    start = time.time()
    result = client.execute_workflow(
        "batch-embeddings",
        {"texts": ["Hello world"]},
        timeout=60,
    )
    elapsed = time.time() - start

    print(f"Workflow ID: {result['id']}")
    print(f"Status: {result['status']}")
    print(f"Time taken: {elapsed:.2f}s")
    print(f"Created at: {result['created_at']}")
    if result.get("completed_at"):
        print(f"Completed at: {result['completed_at']}")


def example_async_execution():
    """Example: Async execution (fire and forget)"""
    print("\n" + "="*50)
    print("Example 7: Async Execution")
    print("="*50)

    client = WorkflowClient()
    result = client.execute_workflow(
        "batch-embeddings",
        {"texts": ["text1", "text2", "text3"]},
        wait=False,
    )

    print(f"Workflow ID: {result['id']}")
    print(f"Status: {result['status']}")
    print(f"Created at: {result['created_at']}")
    print(f"Note: Workflow is running asynchronously. Status is {result['status']}.")


if __name__ == "__main__":
    print("Temporal Workflows API Examples")
    print("================================\n")

    # Run examples (comment out if you don't want to call the actual API)
    try:
        example_batch_embeddings()  # Start with simplest example
        print("\n" + "="*50)
        print("✓ Examples completed successfully!")
        print("="*50)
    except requests.exceptions.ConnectionError:
        print("\n✗ Could not connect to gateway")
        print("Make sure the gateway is running at:", GATEWAY)
    except Exception as e:
        print(f"\n✗ Error: {e}")

    # Show all available methods
    print("\n\nAvailable Methods:")
    print("-" * 50)
    client = WorkflowClient()
    print(f"  - chat_and_embed(model, messages, embed_model)")
    print(f"  - multi_model_chat(models, messages)")
    print(f"  - rag_pipeline(query, documents, model, rerank_model, top_k)")
    print(f"  - batch_embeddings(texts, model)")
    print(f"  - execute_workflow(workflow, input, timeout, wait)")