"""
LLM Mock Service - Ollama-compatible API for AI Studio Proxy

This module provides an Ollama-compatible API that forwards requests
to the main AI Studio Proxy server.
"""

import argparse
import os
import sys
import uuid
import logging
import json
from typing import Dict, Any
from datetime import datetime, UTC

from flask import Flask, request, jsonify
import requests


class FlushingStreamHandler(logging.StreamHandler):
    """Custom log handler that ensures flushing after each emit."""
    def emit(self, record):
        try:
            super().emit(record)
            self.flush()
        except Exception:
            self.handleError(record)


# Configure logging
log_format = '%(asctime)s [%(levelname)s] %(message)s'
formatter = logging.Formatter(log_format)

stderr_handler = FlushingStreamHandler(sys.stderr)
stderr_handler.setFormatter(formatter)
stderr_handler.setLevel(logging.INFO)

root_logger = logging.getLogger()
if root_logger.hasHandlers():
    root_logger.handlers.clear()
root_logger.addHandler(stderr_handler)
root_logger.setLevel(logging.INFO)

logger = logging.getLogger(__name__)

app = Flask(__name__)

# Enabled models configuration
ENABLED_MODELS = {
    "gemini-2.5-pro-preview-05-06",
    "gemini-2.5-flash-preview-04-17",
    "gemini-2.0-flash",
    "gemini-2.0-flash-lite",
    "gemini-1.5-pro",
    "gemini-1.5-flash",
    "gemini-1.5-flash-8b",
}

# API configuration
API_URL = ""
DEFAULT_MAIN_SERVER_PORT = 2048
API_KEY = os.environ.get("API_KEY", "123456")

# Mock responses for fallback
OLLAMA_MOCK_RESPONSES = {
    "What is the capital of France?": "The capital of France is Paris.",
    "Tell me about AI.": "AI is the simulation of human intelligence in machines.",
    "Hello": "Hi! How can I assist you today?"
}


@app.route("/", methods=["GET"])
def root_endpoint():
    """Root endpoint - returns Ollama running status."""
    logger.info("Received root path request")
    return "Ollama is running", 200


@app.route("/api/tags", methods=["GET"])
def tags_endpoint():
    """Ollama /api/tags endpoint - returns enabled models list."""
    logger.info("Received /api/tags request")
    models = []
    for model_name in ENABLED_MODELS:
        family = model_name.split('-')[0].lower() if '-' in model_name else model_name.lower()
        
        if 'llama' in model_name:
            family = 'llama'
            format_type = 'gguf'
            size = 1234567890
            parameter_size = '405B' if '405b' in model_name else 'unknown'
            quantization_level = 'Q4_0'
        elif 'mistral' in model_name:
            family = 'mistral'
            format_type = 'gguf'
            size = 1234567890
            parameter_size = 'unknown'
            quantization_level = 'unknown'
        else:
            format_type = 'unknown'
            size = 9876543210
            parameter_size = 'unknown'
            quantization_level = 'unknown'

        models.append({
            "name": model_name,
            "model": model_name,
            "modified_at": datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%S.%fZ"),
            "size": size,
            "digest": str(uuid.uuid4()),
            "details": {
                "parent_model": "",
                "format": format_type,
                "family": family,
                "families": [family],
                "parameter_size": parameter_size,
                "quantization_level": quantization_level
            }
        })
    logger.info(f"Returning {len(models)} models")
    return jsonify({"models": models}), 200


def generate_ollama_mock_response(prompt: str, model: str) -> Dict[str, Any]:
    """Generate mock Ollama chat response."""
    response_content = OLLAMA_MOCK_RESPONSES.get(
        prompt, f"Echo: {prompt} (Mock response from Ollama service)"
    )
    return {
        "model": model,
        "created_at": datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%SZ"),
        "message": {"role": "assistant", "content": response_content},
        "done": True,
        "total_duration": 123456789,
        "load_duration": 1234567,
        "prompt_eval_count": 10,
        "prompt_eval_duration": 2345678,
        "eval_count": 20,
        "eval_duration": 3456789
    }


def convert_api_to_ollama_response(api_response: Dict[str, Any], model: str) -> Dict[str, Any]:
    """Convert OpenAI format response to Ollama format."""
    try:
        content = api_response["choices"][0]["message"]["content"]
        total_duration = api_response.get("usage", {}).get("total_tokens", 30) * 1000000
        prompt_tokens = api_response.get("usage", {}).get("prompt_tokens", 10)
        completion_tokens = api_response.get("usage", {}).get("completion_tokens", 20)

        return {
            "model": model,
            "created_at": datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%SZ"),
            "message": {"role": "assistant", "content": content},
            "done": True,
            "total_duration": total_duration,
            "load_duration": 1234567,
            "prompt_eval_count": prompt_tokens,
            "prompt_eval_duration": prompt_tokens * 100000,
            "eval_count": completion_tokens,
            "eval_duration": completion_tokens * 100000
        }
    except KeyError as e:
        logger.error(f"Failed to convert API response: missing key {str(e)}")
        return {"error": f"Invalid API response format: missing key {str(e)}"}


@app.route("/api/chat", methods=["POST"])
def ollama_chat_endpoint():
    """Ollama /api/chat endpoint."""
    try:
        data = request.get_json()
        if not data or "messages" not in data:
            return jsonify({"error": "Invalid request: missing 'messages' field"}), 400

        messages = data.get("messages", [])
        if not messages or not isinstance(messages, list):
            return jsonify({"error": "Invalid request: 'messages' must be a non-empty list"}), 400

        model = data.get("model", "llama3.2")
        user_message = next(
            (msg["content"] for msg in reversed(messages) if msg.get("role") == "user"),
            ""
        )
        if not user_message:
            return jsonify({"error": "No user message found"}), 400

        logger.info(f"Processing /api/chat request, model: {model}")

        api_request = {
            "model": model,
            "messages": messages,
            "stream": False,
            "temperature": data.get("temperature", 0.7)
        }
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {API_KEY}"
        }

        try:
            response = requests.post(API_URL, json=api_request, headers=headers, timeout=300000)
            response.raise_for_status()
            api_response = response.json()
            ollama_response = convert_api_to_ollama_response(api_response, model)
            return jsonify(ollama_response), 200
        except requests.RequestException as e:
            logger.error(f"API request failed: {str(e)}")
            response = generate_ollama_mock_response(user_message, model)
            return jsonify(response), 200

    except Exception as e:
        logger.error(f"/api/chat server error: {str(e)}")
        return jsonify({"error": f"Server error: {str(e)}"}), 500


@app.route("/v1/chat/completions", methods=["POST"])
def api_chat_endpoint():
    """OpenAI-compatible /v1/chat/completions endpoint."""
    try:
        data = request.get_json()
        if not data or "messages" not in data:
            return jsonify({"error": "Invalid request: missing 'messages' field"}), 400

        messages = data.get("messages", [])
        if not messages or not isinstance(messages, list):
            return jsonify({"error": "Invalid request: 'messages' must be a non-empty list"}), 400

        model = data.get("model", "grok-3")
        logger.info(f"Processing /v1/chat/completions request, model: {model}")

        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {API_KEY}"
        }

        try:
            response = requests.post(API_URL, json=data, headers=headers, timeout=300000)
            response.raise_for_status()
            api_response = response.json()
            ollama_response = convert_api_to_ollama_response(api_response, model)
            return jsonify(ollama_response), 200
        except requests.RequestException as e:
            logger.error(f"API request failed: {str(e)}")
            return jsonify({"error": f"API request failed: {str(e)}"}), 500

    except Exception as e:
        logger.error(f"/v1/chat/completions server error: {str(e)}")
        return jsonify({"error": f"Server error: {str(e)}"}), 500


def main():
    """Start the LLM mock service."""
    global API_URL

    parser = argparse.ArgumentParser(description="LLM Mock Service for AI Studio Proxy")
    parser.add_argument(
        "--main-server-port",
        type=int,
        default=int(os.environ.get("MAIN_SERVER_PORT", DEFAULT_MAIN_SERVER_PORT)),
        help=f"Port of the main AI Studio Proxy server (default: {DEFAULT_MAIN_SERVER_PORT})"
    )
    parser.add_argument(
        "--port",
        type=int,
        default=11434,
        help="Port for the LLM service (default: 11434)"
    )
    args = parser.parse_args()

    API_URL = f"http://localhost:{args.main_server_port}/v1/chat/completions"
    
    logger.info(f"LLM Mock Service forwarding to: {API_URL}")
    logger.info(f"Starting LLM Mock Service on http://0.0.0.0:{args.port}")
    app.run(host="0.0.0.0", port=args.port, debug=False)


if __name__ == "__main__":
    main()
