from __future__ import annotations

import json

from ..typing import AsyncResult, Messages
from .base_provider import AsyncGeneratorProvider, ProviderModelMixin, AuthFileMixin
from ..requests import StreamSession, get_args_from_nodriver, raise_for_status, merge_cookies
from ..requests import DEFAULT_HEADERS, has_nodriver, has_curl_cffi
from ..providers.response import FinishReason, Usage
from ..errors import ResponseStatusError, ModelNotFoundError
from .. import debug
from .helper import render_messages

def clean_name(name: str) -> str:
    return name.split("/")[-1].replace(
        "-instruct", "").replace(
        "-17b-16e", "").replace(
        "-chat", "").replace(
        "-fp8", "").replace(
        "-fast", "").replace(
        "-int8", "").replace(
        "-awq", "").replace(
        "-qvq", "").replace(
        "-r1", "").replace(
        "meta-llama-", "llama-").replace(
        "-it", "").replace(
        "qwen-", "qwen").replace(
        "qwen", "qwen-")

# models = []
# model_aliases = {clean_name(m.get("name")): m.get("name") for m in models}
# open(__file__, "a").write(f"""# Generated by g4f.models.cloudflare.py
# model_aliases = {model_aliases}
# """)

class Cloudflare(AsyncGeneratorProvider, ProviderModelMixin, AuthFileMixin):
    label = "Cloudflare AI"
    url = "https://playground.ai.cloudflare.com"
    working = False
    use_nodriver = True
    active_by_default = True
    api_endpoint = "https://playground.ai.cloudflare.com/api/inference"
    models_url = "https://playground.ai.cloudflare.com/api/models"
    supports_stream = True
    supports_system_message = True
    supports_message_history = True
    default_model = 'llama-3.3-70b'
    model_aliases = {
        'deepseek-coder-6.7b-base': '@hf/thebloke/deepseek-coder-6.7b-base-awq',
        'deepseek-coder-6.7b': '@hf/thebloke/deepseek-coder-6.7b-instruct-awq',
        'deepseek-math-7b': '@cf/deepseek-ai/deepseek-math-7b-instruct',
        'deepseek-distill-qwen-32b': '@cf/deepseek-ai/deepseek-r1-distill-qwen-32b',
        'discolm-german-7b-v1': '@cf/thebloke/discolm-german-7b-v1-awq',
        'falcon-7b': '@cf/tiiuae/falcon-7b-instruct',
        'gemma-3-12b': '@cf/google/gemma-3-12b-it',
        'gemma-7b': '@hf/google/gemma-7b-it',
        'hermes-2-pro-mistral-7b': '@hf/nousresearch/hermes-2-pro-mistral-7b',
        'llama-2-13b': '@hf/thebloke/llama-2-13b-chat-awq',
        'llama-2-7b-fp16': '@cf/meta/llama-2-7b-chat-fp16',
        'llama-2-7b': '@cf/meta/llama-2-7b-chat-int8',
        'llama-3-8b': '@hf/meta-llama/meta-llama-3-8b-instruct',
        'llama-3.1-8b': '@cf/meta/llama-3.1-8b-instruct-fp8',
        'llama-3.2-11b-vision': '@cf/meta/llama-3.2-11b-vision-instruct',
        'llama-3.2-1b': '@cf/meta/llama-3.2-1b-instruct',
        'llama-3.2-3b': '@cf/meta/llama-3.2-3b-instruct',
        'llama-3.3-70b': '@cf/meta/llama-3.3-70b-instruct-fp8-fast',
        'llama-4-scout': '@cf/meta/llama-4-scout-17b-16e-instruct',
        'llama-guard-3-8b': '@cf/meta/llama-guard-3-8b',
        'llamaguard-7b': '@hf/thebloke/llamaguard-7b-awq',
        'mistral-7b-v0.1': '@hf/thebloke/mistral-7b-instruct-v0.1-awq',
        'mistral-7b-v0.2': '@hf/mistral/mistral-7b-instruct-v0.2',
        'mistral-small-3.1-24b': '@cf/mistralai/mistral-small-3.1-24b-instruct',
        'neural-7b-v3-1': '@hf/thebloke/neural-chat-7b-v3-1-awq',
        'openchat-3.5-0106': '@cf/openchat/openchat-3.5-0106',
        'openhermes-2.5-mistral-7b': '@hf/thebloke/openhermes-2.5-mistral-7b-awq',
        'phi-2': '@cf/microsoft/phi-2',
        'qwen1.5-0.5b': '@cf/qwen/qwen1.5-0.5b-chat',
        'qwen-1.5-1.8b': '@cf/qwen/qwen1.5-1.8b-chat',
        'qwen-1.5-14b': '@cf/qwen/qwen1.5-14b-chat-awq',
        'qwen-1.5-7b': '@cf/qwen/qwen1.5-7b-chat-awq',
        'qwen-2.5-coder-32b': '@cf/qwen/qwen2.5-coder-32b-instruct',
        'qwq-32b': '@cf/qwen/qwq-32b',
        'sqlcoder-7b-2': '@cf/defog/sqlcoder-7b-2',
        'starling-lm-7b-beta': '@hf/nexusflow/starling-lm-7b-beta',
        'tinyllama-1.1b-v1.0': '@cf/tinyllama/tinyllama-1.1b-chat-v1.0',
        'una-cybertron-7b-v2-bf16': '@cf/fblgit/una-cybertron-7b-v2-bf16',
        'zephyr-7b-beta': '@hf/thebloke/zephyr-7b-beta-awq'
    }
    models = list(model_aliases.keys())
    _args: dict = None

    @classmethod
    async def create_async_generator(
        cls,
        model: str,
        messages: Messages,
        proxy: str = None,
        max_tokens: int = 2048,
        **kwargs
    ) -> AsyncResult:
        cache_file = cls.get_cache_file()
        if cls._args is None:
            headers = DEFAULT_HEADERS.copy()
            headers["referer"] = f"{cls.url}"
            headers["origin"] = cls.url
            if cache_file.exists():
                with cache_file.open("r") as f:
                    cls._args = json.load(f)
            elif has_nodriver:
                try:
                    cls._args = await get_args_from_nodriver(cls.url, proxy=proxy)
                except (RuntimeError, FileNotFoundError) as e:
                    debug.log(f"Cloudflare: Nodriver is not available:", e)
                    cls._args = {"headers": headers, "cookies": {}, "impersonate": "chrome"}
            else:
                cls._args = {"headers": headers, "cookies": {}, "impersonate": "chrome"}
        try:
            model = cls.get_model(model)
        except ModelNotFoundError:
            pass
        data = {
            "messages": [{
                **message,
                "parts": [{"type":"text", "text": message["content"]}]} for message in render_messages(messages)],
            "lora": None,
            "model": model,
            "max_tokens": max_tokens,
            "stream": True,
            "system_message":"You are a helpful assistant",
            "tools":[]
        }
        async with StreamSession(**cls._args) as session:
            async with session.post(
                cls.api_endpoint,
                json=data,
            ) as response:
                cls._args["cookies"] = merge_cookies(cls._args["cookies"] , response)
                try:
                    await raise_for_status(response)
                except ResponseStatusError:
                    cls._args = None
                    if cache_file.exists():
                        cache_file.unlink()
                    raise
                async for line in response.iter_lines():
                    if line.startswith(b'0:'):
                        yield json.loads(line[2:])
                    elif line.startswith(b'e:'):
                        finish = json.loads(line[2:])
                        yield Usage(**finish.get("usage"))
                        yield FinishReason(finish.get("finishReason"))
        with cache_file.open("w") as f:
            json.dump(cls._args, f)
