from pydantic import BaseModel, Field

class CategoryResult(BaseModel):
    category: str = Field(description="tech, business, health, or sports")
    confidence: float = Field(ge=0.0, le=1.0)
    reasoning: str

@reasoners_router.reasoner()
async def categorize_text(text: str) -> dict:
    """Categorize text into predefined categories"""
    result = await reasoners_router.ai(
        system="You are a text categorization expert. Categories: tech, business, health, sports.",
        user=f"Categorize this text: {text}",
        schema=CategoryResult
    )
    return result.model_dump()

type CategoryResult struct {
    Category   string  `json:"category"`
    Confidence float64 `json:"confidence"`
    Reasoning  string  `json:"reasoning"`
}

app.RegisterReasoner("categorize_text", func(ctx context.Context, input map[string]any) (any, error) {
    text, _ := input["text"].(string)

    var result CategoryResult
    err := app.AI(ctx, agent.AIRequest{
        System: "You are a text categorization expert. Categories: tech, business, health, sports.",
        User:   "Categorize this text: " + text,
        Schema: &result,
    })

    return result, err
})

@reasoners_router.reasoner()
async def research_topic(topic: str) -> dict:
    """Multi-step research workflow"""

    # Step 1: Generate search queries
    queries = await reasoners_router.ai(
        system="Generate 3 search queries for research.",
        user=f"Topic: {topic}",
        schema=list[str]
    )

    # Step 2: Call search skill (cross-agent communication)
    results = []
    for query in queries:
        result = await app.call(
            "search-agent.web_search",
            query=query
        )
        results.append(result)

    # Step 3: Synthesize findings
    summary = await reasoners_router.ai(
        system="Synthesize research findings into a coherent summary.",
        user=f"Findings: {results}",
        schema=str
    )

    return {
        "topic": topic,
        "queries": queries,
        "summary": summary
    }

app.RegisterReasoner("research_topic", func(ctx context.Context, input map[string]any) (any, error) {
    topic, _ := input["topic"].(string)

    // Step 1: Generate search queries
    var queries []string
    app.AI(ctx, agent.AIRequest{
        System: "Generate 3 search queries for research.",
        User:   "Topic: " + topic,
        Schema: &queries,
    })

    // Step 2: Call search skill
    var results []interface{}
    for _, query := range queries {
        result, _ := app.Call(ctx, "search-agent.web_search", map[string]string{
            "query": query,
        })
        results = append(results, result)
    }

    // Step 3: Synthesize findings
    var summary string
    app.AI(ctx, agent.AIRequest{
        System: "Synthesize research findings.",
        User:   fmt.Sprintf("Findings: %v", results),
        Schema: &summary,
    })

    return map[string]interface{}{
        "topic":   topic,
        "queries": queries,
        "summary": summary,
    }, nil
})

@reasoners_router.reasoner()
async def analyze_with_context(text: str) -> dict:
    """Use multiple agents for comprehensive analysis"""

    # Call sentiment agent
    sentiment = await app.call(
        "sentiment-agent.analyze",
        text=text
    )

    # Call entity extraction agent
    entities = await app.call(
        "nlp-agent.extract_entities",
        text=text
    )

    # Synthesize results
    analysis = await reasoners_router.ai(
        system="Combine sentiment and entity analysis.",
        user=f"Sentiment: {sentiment}, Entities: {entities}",
        schema=dict
    )

    return analysis

app.RegisterReasoner("analyze_with_context", func(ctx context.Context, input map[string]any) (any, error) {
    text, _ := input["text"].(string)

    // Call sentiment agent
    sentiment, _ := app.Call(ctx, "sentiment-agent.analyze", map[string]string{
        "text": text,
    })

    // Call entity extraction agent
    entities, _ := app.Call(ctx, "nlp-agent.extract_entities", map[string]string{
        "text": text,
    })

    // Synthesize results
    var analysis map[string]interface{}
    app.AI(ctx, agent.AIRequest{
        System: "Combine sentiment and entity analysis.",
        User:   fmt.Sprintf("Sentiment: %v, Entities: %v", sentiment, entities),
        Schema: &analysis,
    })

    return analysis, nil
})

@reasoners_router.reasoner()
async def remember_preference(user_id: str, preference: str) -> dict:
    """Store user preference in shared memory"""

    # Store in memory
    await app.memory.set(
        key=f"user:{user_id}:preference",
        value=preference,
        ttl=86400  # 24 hours
    )

    return {"status": "saved", "user_id": user_id}

@reasoners_router.reasoner()
async def get_preference(user_id: str) -> dict:
    """Retrieve user preference"""

    preference = await app.memory.get(
        key=f"user:{user_id}:preference"
    )

    return {"user_id": user_id, "preference": preference}

app.RegisterReasoner("remember_preference", func(ctx context.Context, input map[string]any) (any, error) {
    userID, _ := input["user_id"].(string)
    preference, _ := input["preference"].(string)

    app.Memory.Set(ctx, agent.MemoryKey{
        Key:   "user:" + userID + ":preference",
        Value: preference,
        TTL:   86400,
    })

    return map[string]string{
        "status":  "saved",
        "user_id": userID,
    }, nil
})