# Example agent implementation
# This demonstrates the integration pattern with AI Skills API

import os
import asyncio
import httpx
from typing import List, Dict, Optional

API_URL = os.getenv("API_URL", "http://helm:8675")
API_KEY = os.getenv("API_KEY")

async def get_context(query: str, project: Optional[str] = None) -> Dict:
    """Fetch relevant context from skills API"""
    params = {"query": query}
    if project:
        params["project"] = project
    
    headers = {"X-API-Key": API_KEY} if API_KEY else {}
    
    async with httpx.AsyncClient() as client:
        resp = await client.get(f"{API_URL}/context/rag", params=params, headers=headers)
        resp.raise_for_status()
        return resp.json()

async def compress_messages(messages: List[Dict]) -> Dict:
    """Compress conversation history"""
    headers = {"X-API-Key": API_KEY} if API_KEY else {}
    
    async with httpx.AsyncClient() as client:
        resp = await client.post(f"{API_URL}/compress", json={"messages": messages}, headers=headers)
        resp.raise_for_status()
        return resp.json()

async def store_memory(project: str, key: str, content: str) -> Dict:
    """Store a memory for future reference"""
    headers = {"X-API-Key": API_KEY} if API_KEY else {}
    
    async with httpx.AsyncClient() as client:
        resp = await client.post(
            f"{API_URL}/memory",
            json={"id": key[:8], "project": project, "key": key, "content": content},
            headers=headers
        )
        resp.raise_for_status()
        return resp.json()

async def count_tokens(text: str) -> int:
    """Count tokens using skills API"""
    headers = {"X-API-Key": API_KEY} if API_KEY else {}
    
    async with httpx.AsyncClient() as client:
        resp = await client.get(f"{API_URL}/tokens/count", params={"text": text}, headers=headers)
        resp.raise_for_status()
        return resp.json()["tokens"]

async def chat_loop():
    """Main chat loop - integrate with your LLM of choice"""
    conversation = []
    
    print("Agent ready! Type 'quit' to exit.")
    
    while True:
        user_input = input("\nYou: ")
        if user_input.lower() == 'quit':
            break
        
        # 1. Get relevant context
        context = await get_context(user_input, project="/home/user/projects/myapp")
        context_str = format_context(context)
        
        # 2. Build prompt with context
        system_msg = f"{context_str}\n\nYou are a helpful assistant."
        messages = [{"role": "system", "content": system_msg}]
        messages.extend(conversation[-4:])  # Keep last few turns
        messages.append({"role": "user", "content": user_input})
        
        # 3. Call your LLM here (not included - use OpenAI, Claude, Ollama, etc.)
        # response = await call_llm(messages)
        # For demo, we'll just echo
        response = f"Echo: {user_input}"
        
        # 4. Update conversation
        conversation.append({"role": "user", "content": user_input})
        conversation.append({"role": "assistant", "content": response})
        
        # 5. Compress if getting long
        if len(conversation) > 10:
            compression = await compress_messages(conversation)
            conversation = compression["messages"]
            print(f"\n[Compressed: saved {compression['tokens_saved']} tokens]")
        
        print(f"\nAssistant: {response}")

def format_context(context: Dict) -> str:
    """Format RAG context for inclusion in prompt"""
    parts = []
    
    if context.get("skills"):
        parts.append("## Relevant Skills\n")
        for skill in context["skills"]:
            parts.append(f"### {skill['name']} (relevance: {skill['relevance_score']:.2f})\n{skill['content']}\n")
    
    if context.get("conventions"):
        parts.append("## Project Conventions\n")
        for conv in context["conventions"]:
            parts.append(f"### {conv['name']}\n{conv['content']}\n")
    
    if context.get("snippets"):
        parts.append("## Code Snippets\n")
        for snippet in context["snippets"]:
            parts.append(f"### {snippet['name']} ({snippet['language']})\n```{snippet['language']}\n{snippet['content']}\n```\n")
    
    return "\n".join(parts) if parts else "No relevant context found."

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