import os
from flask import Flask, request, jsonify
import requests
import logging
from dotenv import load_dotenv
import json
from typing import Dict, Any, Tuple, Union

# Load environment variables
load_dotenv()

app = Flask(__name__)

ANTHROPIC_API_URL = os.getenv(
    'ANTHROPIC_API_URL',
    'https://relay.stagwellmarketingcloud.io/anthropic/v1/messages'
)
API_KEY = os.getenv('ANTHROPIC_API_KEY')

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

def transform_to_anthropic_request(
    request_data: Dict[str, Any]
) -> Dict[str, Any]:
    """Transform OpenAI-style request to Anthropic format."""
    if 'messages' in request_data:
        # Chat completion format
        return {
            "model": request_data.get('model', 'claude-3-5-sonnet-20240620'),
            "max_tokens": request_data.get('max_tokens', 1024),
            "messages": request_data['messages']
        }
    else:
        # Regular completion format
        return {
            "model": request_data.get('model', 'claude-3-5-sonnet-20240620'),
            "max_tokens": request_data.get('max_tokens', 1024),
            "messages": [{
                "role": "user",
                "content": request_data.get('prompt', '')
            }]
        }

def make_anthropic_request(
    request_data: Dict[str, Any]
) -> Tuple[Dict[str, Any], int]:
    """Make request to Anthropic API with proper error handling."""
    headers = {
        'Authorization': f'Bearer {API_KEY}',
        'Content-Type': 'application/json',
    }
    
    try:
        anthropic_request = transform_to_anthropic_request(request_data)
        logger.info(
            f"Sending request to Anthropic API: {json.dumps(anthropic_request)}"
        )
        
        response = requests.post(
            ANTHROPIC_API_URL,
            headers=headers,
            json=anthropic_request,
            timeout=30
        )
        response.raise_for_status()
        
        return response.json(), 200
        
    except requests.RequestException as e:
        error_msg = f"Error communicating with Anthropic API: {str(e)}"
        if hasattr(e, 'response') and e.response is not None:
            error_msg += f"\nResponse: {e.response.text}"
        logger.error(error_msg)
        return {"error": error_msg}, e.response.status_code if e.response else 500
    except Exception as e:
        error_msg = f"Unexpected error: {str(e)}"
        logger.error(error_msg)
        return {"error": error_msg}, 500

# Routes remain mostly unchanged, but use new helper functions
@app.route('/v1/completions', methods=['POST'])
def completions():
    """Handle completion-style requests."""
    logger.info(f"Received completion request: {json.dumps(request.json)}")
    
    anthropic_response, status_code = make_anthropic_request(request.json)
    if status_code != 200:
        return jsonify(anthropic_response), status_code
        
    return jsonify({
        "id": "cmpl-" + anthropic_response.get('id', 'default'),
        "object": "text_completion",
        "created": anthropic_response.get('created', 0),
        "choices": [{
            "text": anthropic_response['content'][0]['text'],
            "index": 0,
            "finish_reason": "stop"
        }],
        "usage": {
            "prompt_tokens": -1,
            "completion_tokens": -1,
            "total_tokens": -1
        }
    }), 200

@app.route('/v1/chat/completions', methods=['POST'])
def chat_completions():
    """Handle chat completion-style requests."""
    logger.info(f"Received chat completion request: {json.dumps(request.json)}")
    
    anthropic_response, status_code = make_anthropic_request(request.json)
    if status_code != 200:
        return jsonify(anthropic_response), status_code
        
    return jsonify({
        "id": "chatcmpl-" + anthropic_response.get('id', 'default'),
        "object": "chat.completion",
        "created": anthropic_response.get('created', 0),
        "choices": [{
            "index": 0,
            "message": {
                "role": "assistant",
                "content": anthropic_response['content'][0]['text']
            },
            "finish_reason": "stop"
        }],
        "usage": {
            "prompt_tokens": -1,
            "completion_tokens": -1,
            "total_tokens": -1
        }
    }), 200

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=4224)