import os
import json
import datetime
from pathlib import Path
from dotenv import load_dotenv
from fastmcp import FastMCP, mcp
from pydantic import BaseModel, Field

# Load environment variables
load_dotenv()

# Configuration
DATA_FILE = Path(os.getenv("MEMORY_DATA_FILE", "memory.json"))

# Pydantic model for stored value
class StoredItem(BaseModel):
    key: str = Field(..., description="The key for the item")
    value: str = Field(..., description="The stored value")
    timestamp: str = Field(..., description="ISO 8601 timestamp when stored")

# In-memory cache
_memory_cache: dict[str, StoredItem] = {}

# Load persisted data on startup
def load_data() -> None:
    global _memory_cache
    if DATA_FILE.exists():
        try:
            with DATA_FILE.open("r", encoding="utf-8") as f:
                data = json.load(f)
            for key, item in data.items():
                try:
                    _memory_cache[key] = StoredItem(**item)
                except Exception as exc:
                    print(f"[WARN] Failed to parse item {key}: {exc}")
        except Exception as exc:
            print(f"[ERROR] Unable to read data file {DATA_FILE}: {exc}")

# Persist data to disk
def persist_data() -> None:
    try:
        data = {k: v.dict() for k, v in _memory_cache.items()}
        with DATA_FILE.open("w", encoding="utf-8") as f:
            json.dump(data, f, indent=2)
    except Exception as exc:
        print(f"[ERROR] Unable to write data file {DATA_FILE}: {exc}")

# Tool definitions
@mcp.tool("save")
def save(key: str, value: str) -> str:
    """Save a key-value pair with timestamp."""
    try:
        item = StoredItem(key=key, value=value, timestamp=datetime.datetime.utcnow().isoformat())
        _memory_cache[key] = item
        persist_data()
        return f"Saved key '{key}'."
    except Exception as exc:
        return f"Error saving key '{key}': {exc}"

@mcp.tool("get")
def get(key: str) -> str:
    """Retrieve value for a key."""
    try:
        item = _memory_cache[key]
        return f"{item.value} (stored at {item.timestamp})"
    except KeyError:
        return f"Key '{key}' not found."
    except Exception as exc:
        return f"Error retrieving key '{key}': {exc}"

@mcp.tool("delete")
def delete(key: str) -> str:
    """Delete a key."""
    try:
        del _memory_cache[key]
        persist_data()
        return f"Deleted key '{key}'."
    except KeyError:
        return f"Key '{key}' not found."
    except Exception as exc:
        return f"Error deleting key '{key}': {exc}"

@mcp.tool("list_keys")
def list_keys() -> str:
    """List all keys."""
    try:
        keys = ", ".join(_memory_cache.keys()) or "No keys"
        return f"Keys: {keys}"
    except Exception as exc:
        return f"Error listing keys: {exc}"

# Namespace tools
@mcp.tool("save_with_namespace")
def save_with_namespace(namespace: str, key: str, value: str) -> str:
    """Save a key-value pair under a namespace."""
    namespaced_key = f"{namespace}:{key}"
    return save(namespaced_key, value)

@mcp.tool("get_by_namespace")
def get_by_namespace(namespace: str, key: str) -> str:
    """Retrieve a value from a namespace."""
    namespaced_key = f"{namespace}:{key}"
    return get(namespaced_key)

# Initialize server
mcpServer = FastMCP()
# Register tools automatically via decorators

# Load existing data
load_data()

# Demo client if invoked with --client
if __name__ == "__main__":
    import sys

    if len(sys.argv) > 1 and sys.argv[1] == "--client":
        # Simple interactive client using stdio transport
        print("Running client demo. Type 'exit' to quit.")
        print("Available tools: save, get, delete, list_keys, save_with_namespace, get_by_namespace")
        while True:
            try:
                user_input = input("> ")
            except EOFError:
                break
            if user_input.strip().lower() == "exit":
                break
            # Send command via mcp.run() with input string
            try:
                response = mcp.run(user_input)
                print(response)
            except Exception as e:
                print(f"Error: {e}")
    else:
        # Run server with stdio transport
        print("Starting FastMCP server with stdio transport...")
        mcp.run()