import os
import urllib.request
import uvicorn
from llama_cpp.server.app import create_app, Settings

from starlette_context.plugins import RequestIdPlugin

def patch_request_id_plugin():
    original_init = RequestIdPlugin.__init__
    def my_init(self, *args, **kwargs):
        kwargs['validate'] = False
        original_init(self, *args, **kwargs)
    RequestIdPlugin.__init__ = my_init

patch_request_id_plugin()

MODEL_PATH = "/model.gguf"
MODEL_URL = "https://huggingface.co/bartowski/Meta-Llama-3.1-8B-Instruct-GGUF/resolve/main/Meta-Llama-3.1-8B-Instruct-Q4_K_M.gguf"
EXPECTED_SIZE = 4920000000  # ~4.92 GB

def download_model():
    if os.path.exists(MODEL_PATH):
        size = os.path.getsize(MODEL_PATH)
        if size >= EXPECTED_SIZE:  # Ajustado para maior ou igual
            print(f"Modelo já existe ({size / 1e9:.2f} GB). Pulando download.")
            return
        else:
            print(f"Arquivo incompleto ({size / 1e9:.2f} GB). Removendo e baixando novamente.")
            os.remove(MODEL_PATH)

    print("Baixando modelo Q8_0 (~8.5 GB)...")
    attempt = 0
    while attempt < 3:
        try:
            urllib.request.urlretrieve(MODEL_URL, MODEL_PATH)
            size = os.path.getsize(MODEL_PATH)
            if size >= EXPECTED_SIZE:
                print(f"Download concluído: {size / 1e9:.2f} GB")
                return
            else:
                print(f"Download incompleto ({size / 1e9:.2f} GB). Tentativa {attempt + 1}/3")
                if os.path.exists(MODEL_PATH):
                    os.remove(MODEL_PATH)
        except Exception as e:
            print(f"Erro no download: {e}. Tentativa {attempt + 1}/3")
            if os.path.exists(MODEL_PATH):
                os.remove(MODEL_PATH)
        attempt += 1

    raise RuntimeError("Falha ao baixar o modelo após 3 tentativas.")

download_model()

settings = Settings(
    model=MODEL_PATH,
    model_alias="julia-llm",
    n_ctx=8192,
    n_threads=2,
    n_gpu_layers=0,
    chat_format="llama-3",
    host="0.0.0.0",
    port=7860,
)

app = create_app(settings=settings)


if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860, http="h11")