from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import Any
import uvicorn
import json
import redis
import os

# Initialize Redis connection
def get_redis_client():
    redis_host = os.getenv('REDIS_HOST', 'localhost')
    redis_port = int(os.getenv('REDIS_PORT', 6379))
    redis_db = int(os.getenv('REDIS_DB', 0))

    try:
        client = redis.StrictRedis(
            host=redis_host,
            port=redis_port,
            db=redis_db,
            decode_responses=True
        )
        # Test connection
        client.ping()
        return client
    except redis.ConnectionError as e:
        raise HTTPException(status_code=503, detail=f"Could not connect to Redis: {str(e)}")

redis_client = get_redis_client()

app = FastAPI(
    title="Remember Service",
    description="An OpenAPI-compatible service that allows an LLM to save details by remembering them.",
    version="1.0.0"
)

class RememberRequest(BaseModel):
    """
    The request model for the remember tool.
    """
    content: Any
    metadata: dict = {}

@app.post("/remember",
          summary="Store information to remember",
          description="Saves the provided content and optional metadata in Redis.")
def remember_endpoint(request: RememberRequest):
    """
    Endpoint to remember information.
    Stores the received payload in Redis and prints it to stdout in JSON format.
    """
    payload = {
        "content": request.content,
        "metadata": request.metadata
    }
    payload_json = json.dumps(payload, ensure_ascii=False)

    # Generate a unique key based on timestamp
    import time
    key = f"remember:{int(time.time())}"

    # Store in Redis
    redis_client.set(key, payload_json)

    print(json.dumps(payload, indent=2, ensure_ascii=False))
    return {"status": "remembered", "key": key}

@app.get("/recall",
         summary="Recall information previously remembered",
         description="Searches for content in Redis based on query string. Returns matching entries.")
def recall_endpoint(query: str):
    """
    Endpoint to recall information from Redis.
    Searches for entries containing the query string in their content or metadata.
    """
    try:
        # Get all keys matching the remember pattern
        keys = redis_client.keys("remember:*")

        results = []
        for key in keys:
            data = redis_client.get(key)
            if data:
                payload = json.loads(data)
                # Check if query is in content or metadata
                if (isinstance(payload.get("content"), str) and query.lower() in payload["content"].lower()) or \
                   (isinstance(payload.get("metadata"), dict) and any(
                       isinstance(v, str) and query.lower() in v.lower()
                       for v in payload["metadata"].values()
                   )):
                    results.append({
                        "key": key,
                        "content": payload["content"],
                        "metadata": payload["metadata"]
                    })

        return {"results": results, "count": len(results)}

    except redis.RedisError as e:
        raise HTTPException(status_code=500, detail=f"Error searching Redis: {str(e)}")

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)