from flask import Flask, request, jsonify
from datetime import datetime
import os
import json
from threading import Lock
from openai import OpenAI
import logging
from flask_swagger_ui import get_swaggerui_blueprint

# Configure logging
DATA_DIR = "data"
LOG_FILE = os.path.join(DATA_DIR, "slop_with_models.log")
logging.basicConfig(
    level=logging.DEBUG,
    format="%(asctime)s - %(levelname)s - %(message)s",
    handlers=[logging.FileHandler(LOG_FILE), logging.StreamHandler()]
)
logger = logging.getLogger(__name__)

app = Flask(__name__)

# Swagger UI setup
SWAGGER_URL = "/openapi"
API_URL = "/static/openapi.yaml"
swaggerui_blueprint = get_swaggerui_blueprint(
    SWAGGER_URL, API_URL, config={"app_name": "SLOP API"}
)
app.register_blueprint(swaggerui_blueprint, url_prefix=SWAGGER_URL)

# Global model-to-client map
MODEL_CLIENT_MAP = {}

# Memory setup
MEMORY_FILE = os.path.join(DATA_DIR, "memory.json")
memory_lock = Lock()

# Ensure data directory exists
if not os.path.exists(DATA_DIR):
    os.makedirs(DATA_DIR)
    logger.info(f"Created data directory: {DATA_DIR}")

# File-based memory functions
def load_memory_from_file():
    try:
        if os.path.exists(MEMORY_FILE):
            with open(MEMORY_FILE, "r") as f:
                data = json.load(f)
                logger.debug(f"Loaded memory from {MEMORY_FILE}: {data}")
                return data
        logger.debug(f"No memory file found at {MEMORY_FILE}, returning empty dict")
        return {}
    except Exception as e:
        logger.error(f"Error loading memory from {MEMORY_FILE}: {str(e)}", exc_info=True)
        return {}

def save_memory_to_file(memory_data):
    try:
        with open(MEMORY_FILE, "w") as f:
            json.dump(memory_data, f)
        logger.debug(f"Saved memory to {MEMORY_FILE}: {memory_data}")
    except Exception as e:
        logger.error(f"Error saving memory to {MEMORY_FILE}: {str(e)}", exc_info=True)

# Load endpoints
ENDPOINTS = []
for i in range(1000):
    endpoint = os.getenv(f"MODEL_ENDPOINT_{i}")
    if endpoint:
        api_key = os.getenv(f"MODEL_API_KEY_{i}", "not-needed")
        logger.info(f"Loaded endpoint {i}: {endpoint} with API key: {api_key[:4]}...")
        ENDPOINTS.append(
            {
                "name": f"endpoint_{i}",
                "base_url": endpoint,
                "api_key": api_key,
            }
        )
    else:
        logger.debug(f"No endpoint found for MODEL_ENDPOINT_{i}")

def initialize_model_map():
    logger.info("Starting model map initialization")
    MODEL_CLIENT_MAP.clear()
    if not ENDPOINTS:
        logger.warning("No endpoints configured to initialize models")
        return
    for ep in ENDPOINTS:
        base_url = ep["base_url"]
        api_key = ep["api_key"]
        endpoint_name = ep["name"]
        logger.info(f"Initializing client for {endpoint_name} at {base_url}")
        client = OpenAI(base_url=base_url, api_key=api_key)
        try:
            logger.debug(f"Attempting to list models from {endpoint_name}")
            response = client.models.list()
            model_list = response.data
            logger.debug(f"Models retrieved from {endpoint_name}: {[m.id for m in model_list]}")
        except Exception as e:
            logger.error(f"Failed to list models for {endpoint_name}: {str(e)}", exc_info=True)
            continue
        for m in model_list:
            model_id = m.id
            if model_id:
                if model_id in MODEL_CLIENT_MAP:
                    logger.warning(f"Duplicate model ID '{model_id}' found at {endpoint_name}")
                else:
                    MODEL_CLIENT_MAP[model_id] = client
                    logger.info(f"Registered model '{model_id}' from {endpoint_name}")
            else:
                logger.warning(f"Model with no ID encountered from {endpoint_name}")
    logger.info(f"Model map initialized with models: {list(MODEL_CLIENT_MAP.keys())}")

# SLOP components
tools = {
    "calculator": {
        "id": "calculator",
        "description": "Basic math",
        "execute": lambda params: {"result": eval(params["expression"])},
    },
    "greet": {
        "id": "greet",
        "description": "Says hello",
        "execute": lambda params: {"result": f"Hello, {params['name']}!"},
    },
}
resources = {
    "hello": {"id": "hello", "content": "Hello, SLOP!"},
    "foo/bar": {"id": "foo/bar", "content": "Nested Foo Bar"},
    "foo/baz": {"id": "foo/baz", "content": "Nested Foo Baz"},
}

# Memory endpoints with lock and timeout
LOCK_TIMEOUT = 2  # Seconds to wait for lock acquisition

@app.route("/memory", methods=["POST"])
def store_memory():
    logger.info("Received /memory POST request")
    data = request.json
    logger.debug(f"Store memory data: {data}")
    if not data or "key" not in data or "value" not in data:
        logger.warning(f"Invalid memory store request: {data}")
        return jsonify({"error": "Missing 'key' or 'value'"}), 400
    key, value = data["key"], data["value"]
    if memory_lock.acquire(timeout=LOCK_TIMEOUT):
        try:
            memory = load_memory_from_file()  # Always load fresh state
            memory[key] = value
            save_memory_to_file(memory)
            logger.info(f"Stored in memory: {key} = {value}")
            logger.debug(f"Current memory state: {memory}")
            return jsonify({"status": "stored"}), 200
        finally:
            memory_lock.release()
    else:
        logger.error(f"Failed to acquire lock for storing {key} within {LOCK_TIMEOUT} seconds")
        return jsonify({"error": "Memory lock timeout"}), 503

@app.route("/memory/<key>", methods=["GET"])
def get_memory(key):
    logger.info(f"Received /memory/{key} GET request")
    if memory_lock.acquire(timeout=LOCK_TIMEOUT):
        try:
            memory = load_memory_from_file()  # Always load fresh state
            value = memory.get(key)
            logger.debug(f"Retrieved memory for {key}: {value}")
            if value is None:
                logger.warning(f"Key not found in memory: {key}")
            return jsonify({"value": value}), 200
        finally:
            memory_lock.release()
    else:
        logger.error(f"Failed to acquire lock for retrieving {key} within {LOCK_TIMEOUT} seconds")
        return jsonify({"error": "Memory lock timeout"}), 503

@app.route("/memory", methods=["GET"])
def list_memory():
    logger.info("Received /memory GET request")
    if memory_lock.acquire(timeout=LOCK_TIMEOUT):
        try:
            memory = load_memory_from_file()  # Always load fresh state
            keys = list(memory.keys())
            memory_state = memory.copy()
            logger.debug(f"Memory keys: {keys}")
            logger.debug(f"Full memory state: {memory_state}")
            return jsonify({"keys": keys}), 200
        finally:
            memory_lock.release()
    else:
        logger.error(f"Failed to acquire lock for listing memory within {LOCK_TIMEOUT} seconds")
        return jsonify({"error": "Memory lock timeout"}), 503

@app.route("/memory/<key>", methods=["DELETE"])
def delete_memory(key):
    logger.info(f"Received /memory/{key} DELETE request")
    if memory_lock.acquire(timeout=LOCK_TIMEOUT):
        try:
            memory = load_memory_from_file()  # Always load fresh state
            if key not in memory:
                logger.warning(f"Key not found for deletion: {key}")
                return jsonify({"error": "Key not found"}), 404
            old_value = memory[key]
            del memory[key]
            save_memory_to_file(memory)
            logger.info(f"Deleted from memory: {key} (was {old_value})")
            logger.debug(f"Updated memory state: {memory}")
            return jsonify({"status": "deleted"}), 200
        finally:
            memory_lock.release()
    else:
        logger.error(f"Failed to acquire lock for deleting {key} within {LOCK_TIMEOUT} seconds")
        return jsonify({"error": "Memory lock timeout"}), 503

# Other endpoints (unchanged for brevity)
@app.route("/chat", methods=["POST"])
def chat():
    logger.info("Received /chat request")
    data = request.json
    logger.debug(f"Request data: {data}")
    message = data["messages"][0]["content"] if data.get("messages") else "nothing"
    model_id = data.get("model") or (
        list(MODEL_CLIENT_MAP.keys())[0] if MODEL_CLIENT_MAP else None
    )
    logger.info(f"Selected model: {model_id}, message: {message}")
    if not model_id or model_id not in MODEL_CLIENT_MAP:
        logger.error(f"Invalid or missing model_id: {model_id}")
        return jsonify({"error": "Model not found"}), 404
    client = MODEL_CLIENT_MAP[model_id]
    try:
        response = client.chat.completions.create(
            model=model_id,
            messages=[
                {"role": m["role"], "content": m["content"]}
                for m in data.get("messages", [])
            ]
            or [{"role": "user", "content": message}],
        )
        response_content = response.choices[0].message.content
        logger.debug(f"Chat response from {model_id}: {response_content}")
        return (
            jsonify(
                {
                    "choices": [
                        {"message": {"content": response_content}}
                    ]
                }
            ),
            200,
        )
    except Exception as e:
        logger.error(f"Chat error with model {model_id}: {str(e)}", exc_info=True)
        return jsonify({"error": str(e)}), 500

@app.route("/models", methods=["GET"])
def list_models():
    logger.info("Received /models request")
    models = list(MODEL_CLIENT_MAP.keys())
    logger.debug(f"Returning models: {models}")
    return jsonify({"models": models}), 200

@app.route("/tools", methods=["GET"])
def list_tools():
    logger.info("Received /tools request")
    tool_list = [{"id": k, "description": v["description"]} for k, v in tools.items()]
    logger.debug(f"Returning tools: {tool_list}")
    return jsonify({"tools": tool_list}), 200

@app.route("/tools/<tool_id>", methods=["POST"])
def use_tool(tool_id):
    logger.info(f"Received /tools/{tool_id} request")
    if tool_id not in tools:
        logger.error(f"Tool not found: {tool_id}")
        return jsonify({"error": "Tool not found"}), 404
    data = request.json or {}
    logger.debug(f"Tool {tool_id} input data: {data}")
    if tool_id == "calculator" and "expression" not in data:
        logger.warning("Missing 'expression' for calculator tool")
        return jsonify({"error": "Missing 'expression'"}), 400
    if tool_id == "greet" and "name" not in data:
        logger.warning("Missing 'name' for greet tool")
        return jsonify({"error": "Missing 'name'"}), 400
    try:
        result = tools[tool_id]["execute"](data)
        logger.debug(f"Tool {tool_id} result: {result}")
        return jsonify(result), 200
    except Exception as e:
        logger.error(f"Error executing tool {tool_id}: {str(e)}", exc_info=True)
        return jsonify({"error": str(e)}), 500

@app.route("/resources", methods=["GET"])
def list_resources():
    logger.info("Received /resources request")
    resource_list = list(resources.values())
    logger.debug(f"Returning resources: {resource_list}")
    return jsonify({"resources": resource_list}), 200

@app.route("/resources/<resource_id>", methods=["GET"])
def get_resource(resource_id):
    logger.info(f"Received /resources/{resource_id} GET request")
    if resource_id in resources:
        resource = resources[resource_id]
        logger.debug(f"Exact match found for {resource_id}: {resource}")
        return jsonify(resource), 200
    else:
        # Prefix search for nested resources
        prefix = f"{resource_id}/"
        matching_resources = [
            res for key, res in resources.items() if key.startswith(prefix)
        ]
        if matching_resources:
            logger.debug(f"Prefix search for {resource_id} found: {matching_resources}")
            return jsonify({"resources": matching_resources}), 200
        logger.warning(f"No resource or nested resources found for {resource_id}")
        return jsonify({"error": "Resource not found"}), 404

@app.route("/resources/<resource_id>", methods=["PUT"])
def update_resource(resource_id):
    logger.info(f"Received /resources/{resource_id} PUT request")
    data = request.json
    logger.debug(f"Update resource data: {data}")
    if not data or "content" not in data:
        logger.warning(f"Invalid resource update request: {data}")
        return jsonify({"error": "Missing 'content'"}), 400
    # Update or create the resource
    resources[resource_id] = {
        "id": resource_id,
        "content": data["content"]
    }
    logger.info(f"Updated/created resource {resource_id}: {resources[resource_id]}")
    return jsonify({"status": "updated", "resource": resources[resource_id]}), 200

@app.route("/pay", methods=["POST"])
def pay():
    logger.info("Received /pay request")
    data = request.json
    logger.debug(f"Pay request data: {data}")
    amount = data.get("amount", 0)
    transaction_id = f"tx_{int(datetime.now().timestamp())}"
    logger.info(f"Processed payment of {amount} with transaction ID: {transaction_id}")
    return (
        jsonify(
            {
                "transaction_id": transaction_id,
                "status": "success",
            }
        ),
        200,
    )

# Initialize model map on startup
logger.info("Starting application initialization")
initialize_model_map()
logger.info("Application initialization completed")

if __name__ == "__main__":
    logger.info("Starting Flask application on port 31337")
    app.run(debug=True, port=31337)
    logger.info("Flask application stopped")