import logging
import sys
import os
import json
from flask import Flask, request, jsonify
from llama_cpp import Llama
from huggingface_hub import hf_hub_download

# Inisialisasi Flask app
app = Flask(__name__)

# Konfigurasi Logging
logging.basicConfig(stream=sys.stderr, level=logging.INFO)
logger = logging.getLogger(__name__)

# Set environment variables untuk nonaktifkan cache
os.environ["TRANSFORMERS_OFFLINE"] = "1"
os.environ["HF_DATASETS_OFFLINE"] = "1"
os.environ["HUGGINGFACE_HUB_CACHE"] = "/tmp"  # Arahkan cache ke /tmp
os.environ["HFC_USER_AGENT_DISABLE_TELEMETRY"] = "1"  # Nonaktifkan telemetry

# Global Variables
MODEL_NAME = "second-state/Qwen3-0.6B-GGUF"
MODEL_BASENAME = "Qwen3-0.6B-Q4_K_S.gguf"
MODEL_PATH = "/app/models"
MODEL_FILE = os.path.join(MODEL_PATH, MODEL_BASENAME)

# Load Model (Qwen3-0.6B-GGUF)
def load_model():
    logger.info("Sedang memuat model GGUF dari lokal...")
    try:
        # Pastikan model ada menggunakan huggingface_hub
        logger.info(f"Memastikan model ada dengan hf_hub_download: {MODEL_NAME}, {MODEL_BASENAME}")
        hf_hub_download(repo_id=MODEL_NAME, filename=MODEL_BASENAME, local_dir=MODEL_PATH, local_dir_use_symlinks=False)
        logger.info("Model berhasil diunduh menggunakan hf_hub_download")
        
        # Periksa apakah file model ada
        if not os.path.exists(MODEL_FILE):
            logger.error(f"File model tidak ditemukan: {MODEL_FILE}")
            logger.error(f"Cek isi direktori /app/models: {os.listdir('/app/models') if os.path.exists('/app/models') else 'Direktori tidak ditemukan'}")
            return None

        logger.info(f"Memuat model dari path: {MODEL_FILE}")
        llm = Llama(
            model_path=MODEL_FILE,
            n_gpu_layers=0,  # Jalankan di CPU
            n_threads=4,
            verbose=True,
            n_ctx=1024,
        )
        logger.info("Model GGUF berhasil dimuat!")
        return llm
    except Exception as e:
        logger.error(f"Error loading model GGUF: {e}", exc_info=True)
        return None

llm = load_model()

def ask_ai(prompt, llm_model):
    try:
        if llm_model is None:
            return "Model gagal dimuat. Periksa log untuk detailnya."

        # Format prompt untuk Qwen
        formatted_prompt = f"Human: {prompt}\n<|file_separator|>Assistant:"

        # Jalankan model
        output = llm_model(
            prompt, # formatted_prompt di model gemma ini tidak jalan
            max_tokens=512,
            temperature=0.7,
            top_p=0.9,
            stop=["<|file_separator|>"], #stop untuk Qwen
            echo=False,
            stream=False #Non Stream
        )

        # Extract text jawaban
        jawaban = output['choices'][0]['text']

        return jawaban.strip()
    except Exception as e:
        logger.error(f"Error selama inferensi: {e}", exc_info=True)
        return f"Error: {e}"

# Endpoint / (untuk menerima POST request)
@app.route('/', methods=['POST'])
def ask_ai_endpoint():
    try:
        data = request.get_json()  # Ambil data JSON dari request body
        prompt = data.get('prompt')  # Ambil prompt dari data JSON
        if not prompt:
            return jsonify({"error": "Prompt tidak ditemukan dalam request body"}), 400

        jawaban = ask_ai(prompt, llm)  # Dapatkan jawaban dari model
        return jsonify({"jawaban": jawaban})  # Kembalikan jawaban sebagai JSON
    except Exception as e:
        logger.error(f"Error pada endpoint /: {e}", exc_info=True)
        return jsonify({"error": str(e)}), 500

@app.route('/health', methods=['GET'])  # Tambahkan health check
def health_check():
    return jsonify({"status": "healthy"})

# Jalankan aplikasi Flask jika dijalankan langsung (untuk testing lokal)
if __name__ == '__main__':
    import os
    port = int(os.environ.get("PORT", 8501))
    app.run(debug=True, host='0.0.0.0', port=port)