import os
import requests
import json
import logging
from flask import Flask, request, jsonify
from ddtrace import patch_all, tracer
from ddtrace.llmobs import LLMObs
from dotenv import load_dotenv

# Load environment variables from .env file if it exists
load_dotenv()

# Set up logging for debugging
logging.basicConfig(level=logging.DEBUG)

# Enable Datadog tracing
patch_all()

# Enable LLM observability
LLMObs.enable()

app = Flask(__name__)

# Ensure environment variables are set
os.environ['DD_LLMOBS_ENABLED'] = '1'
os.environ['DD_LLMOBS_ML_APP'] = 'anything-api'
os.environ['DD_LLMOBS_AGENTLESS_ENABLED'] = '1'

@app.route('/llm_call', methods=['POST'])
def handle_llm_call():
    logging.debug("Received a request at /llm_call")

    # Extract data from the incoming request (e.g., message to LLM)
    data = request.get_json()
    message = data.get("message")

    if not message:
        logging.error("No message provided in the request.")
        return jsonify({"error": "No message provided"}), 400

    url = 'https://severian-anything.hf.space/api/v1/workspace/scoreboard/chat'
    headers = {
        'accept': 'application/json',
        'Authorization': 'Bearer TYQYM46-RPCMQ98-GCGJMNB-Q23K6HC',  # Replace with actual token securely
        'Content-Type': 'application/json'
    }
    payload = {
        "message": message,
        "mode": "query"
    }

    # Use LLMObs to trace the LLM API call
    with LLMObs.trace("llm_api_call", service="anything-api", resource="chat") as span:
        logging.debug("Starting trace for LLM API call.")
        span.set_tag("llm.request.model", "anything-api")
        span.set_tag("llm.request.input", message)

        try:
            # Make the actual call to the LLM API
            response = requests.post(url, headers=headers, data=json.dumps(payload))
            response.raise_for_status()
            response_data = response.json()
            bot_response = response_data.get("textResponse")

            # Set LLM-specific tags and metrics
            span.set_tag("llm.response.output", bot_response)
            span.set_metric("llm.tokens.prompt", len(message.split()))
            span.set_metric("llm.tokens.completion", len(bot_response.split()))
            span.set_metric("llm.tokens.total", len(message.split()) + len(bot_response.split()))

            logging.debug(f"LLM API response: {bot_response}")

            return jsonify({"bot_response": bot_response})

        except requests.RequestException as e:
            span.set_tag("error", True)
            span.set_tag("error.msg", str(e))
            logging.error(f"Request failed: {e}")
            return jsonify({"error": f"Request failed: {e}"}), 500

        except Exception as e:
            span.set_tag("error", True)
            span.set_tag("error.msg", str(e))
            logging.error(f"An error occurred: {e}")
            return jsonify({"error": f"An error occurred: {e}"}), 500

if __name__ == "__main__":
    logging.info("Starting Flask app on port 7860")
    app.run(host='0.0.0.0', port=7860)