#!/usr/bin/env python3
"""
Local interactive chat client for Atlas API
Connects to the FastAPI server running on localhost:7860
"""

import httpx
import asyncio
import sys
import json

API_BASE_URL = "http://localhost:7860"

# Global session tracking
_session_id = None

async def send_chat_message(prompt: str, use_search: bool = True) -> dict:
    """Send a chat message to the API and return the response"""
    global _session_id
    
    async with httpx.AsyncClient(timeout=30.0) as client:
        try:
            # Prepare headers with session ID if available
            headers = {
                "User-Agent": "local_app.py/1.0",
                "Content-Type": "application/json"
            }
            if _session_id:
                headers["X-Session-ID"] = _session_id
            
            response = await client.post(
                f"{API_BASE_URL}/chat",
                json={
                    "prompt": prompt,
                    "max_new_tokens": 500,
                    "use_search": use_search,
                    "temperature": 0.7
                },
                headers=headers
            )
            response.raise_for_status()
            
            # Extract session ID from response headers if available
            if "X-Session-ID" in response.headers:
                _session_id = response.headers["X-Session-ID"]
            
            return response.json()
        except httpx.ConnectError:
            return {"error": "Could not connect to API server. Make sure app.py is running on localhost:7860"}
        except httpx.TimeoutException:
            return {"error": "Request timed out. Please try again."}
        except httpx.HTTPStatusError as e:
            return {"error": f"API error: {e.response.status_code} - {e.response.text}"}
        except Exception as e:
            return {"error": f"Unexpected error: {str(e)}"}

async def main():
    """Main chat loop"""
    global _session_id
    
    print("Atlas Chat Client with Analytics")
    print("Type your messages and press Enter. Empty input or Ctrl+C to exit.")
    print("Type 'search:' before your message to enable web search.")
    print("-" * 50)
    
    try:
        while True:
            try:
                # Get user input
                user_input = input("\nYou: ").strip()
                
                # Exit on empty input
                if not user_input:
                    print("Goodbye!")
                    break
                
                # Check for search prefix
                use_search = True
                if user_input.startswith("search:"):
                    use_search = True
                    user_input = user_input[7:].strip()
                    if not user_input:
                        print("Please provide a message after 'search:'")
                        continue
                
                # Send message to API
                print("AI: ", end="", flush=True)
                response = await send_chat_message(user_input, use_search)
                
                # Handle response
                if "error" in response:
                    print(f"Error: {response['error']}")
                else:
                    print(response.get("response", "No response received"))
                    
                    # Show search results if available
                    if response.get("search_results"):
                        print(f"\n[Found {len(response['search_results'])} search results]")
                    
                    # Show session info on first message
                    if _session_id:
                        print(f"\n[Session: {_session_id[:8]}...]")
                        
            except KeyboardInterrupt:
                print("\nGoodbye!")
                break
            except EOFError:
                print("\nGoodbye!")
                break
                
    except Exception as e:
        print(f"Fatal error: {e}")
        sys.exit(1)

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