from flask import Flask, render_template, request, jsonify, Response
from duckduckgo_search import DDGS
import json
import time
import sys
import argparse
import logging

# Configure logging
logging.basicConfig(level=logging.INFO, 
                    format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

app = Flask(__name__)

class DDGChatAPI:
    @staticmethod
    def respond_json(response, key="message"):
        """Create an OpenAI-compatible response format."""
        return {
            "id": f"chatcmpl-{int(time.time())}",
            "object": "chat.completion",
            "created": int(time.time()),
            "choices": [{
                "index": 0,
                key: response,
                "finish_reason": "stop"
            }],
            "usage": {
                "prompt_tokens": 0,
                "completion_tokens": 0,
                "total_tokens": 0
            }
        }

    @staticmethod
    def validate_messages(messages):
        """Validate message format."""
        if not isinstance(messages, list):
            return {
                "error": {
                    "message": "'messages' must be a list",
                    "code": "invalid_message_list"
                }
            }, 400

        for message in messages:
            if not isinstance(message, dict) or 'role' not in message or 'content' not in message:
                return {
                    "error": {
                        "message": "Each message must have a 'role' and a 'content'",
                        "code": "invalid_message"
                    }
                }, 400

        return None, None

    @staticmethod
    def wait_for_full_response(last_message, model='gpt-4o-mini', timeout=120):
        """
        Wait for a complete response from DuckDuckGo chat
        
        Args:
            last_message (str): User's message
            model (str): Selected AI model
            timeout (int): Maximum wait time in seconds
        
        Returns:
            str: Full AI response
        """
        start_time = time.time()
        ddgs = DDGS()
        
        attempts = 0
        max_attempts = 5
        
        while attempts < max_attempts:
            try:
                # Attempt to get response
                response = ddgs.chat(last_message, model=model)
                
                # Clean and validate response
                cleaned_response = response.strip()
                
                # Check response quality
                if len(cleaned_response) >= 50:
                    logger.info(f"Successfully generated response in {attempts + 1} attempt(s)")
                    return cleaned_response
                
                # If response is too short, wait and retry
                attempts += 1
                time.sleep(2)  # Wait between attempts
                
                # Break if total timeout is exceeded
                if time.time() - start_time > timeout:
                    break
            
            except Exception as e:
                logger.error(f"Attempt {attempts + 1} failed: {str(e)}")
                attempts += 1
                time.sleep(2)
        
        # Fallback message if no good response generated
        return "I apologize, but I'm unable to generate a complete response at the moment. Please try again later."

@app.route("/")
def homepage():
    return render_template("index.html")

@app.route("/v1/chat/completions", methods=["POST"])
def chat_completions():
    # Get request data
    data = request.json
    messages = data.get("messages", [])
    model = data.get("model", "gpt-4o-mini")
    stream = data.get("stream", False)
    timeout = data.get("timeout", 120)

    # Validate messages
    errors, status = DDGChatAPI.validate_messages(messages)
    if errors:
        return jsonify(errors), status

    # Extract the last message content
    last_message = messages[-1]["content"]

    # Log incoming request
    logger.info(f"Received chat request: model={model}, message={last_message[:100]}...")

    try:
        # Get full response
        response = DDGChatAPI.wait_for_full_response(last_message, model, timeout)
        
        # Log response generation
        logger.info(f"Response generated (length: {len(response)} chars)")

        # Handle streaming response
        if stream:
            def generate():
                # Split response into chunks for streaming
                max_chunk_size = 50
                chunks = [response[i:i+max_chunk_size] for i in range(0, len(response), max_chunk_size)]
                
                for chunk in chunks:
                    delta_response = {
                        "role": "assistant",
                        "content": chunk
                    }
                    yield f"data: {json.dumps(DDGChatAPI.respond_json(delta_response, 'delta'))}\n\n"
                
                yield "data: [DONE]\n\n"

            return Response(generate(), mimetype='text/event-stream')

        # Regular JSON response
        response_data = {
            "role": "assistant",
            "content": response
        }
        return jsonify(DDGChatAPI.respond_json(response_data))

    except Exception as e:
        logger.error(f"Error processing request: {str(e)}")
        return jsonify({
            "error": {
                "message": str(e),
                "code": "ddg_chat_error"
            }
        }), 500

@app.route("/v1/models", methods=["GET"])
def list_models():
    """Provide a list of available models."""
    return jsonify({
        "data": [
            {"id": "gpt-4o-mini"},
            {"id": "claude-3-haiku"},
            {"id": "llama-3.1-70b"},
            {"id": "mixtral-8x7b"}
        ],
        "object": "list"
    })

def main():
    app.run(host='0.0.0.0', port=7860)

if __name__ == "__main__":
    main()