# ai_server.py
# سرور Flask برای سرویس دهی به مدل هوش مصنوعی

from flask import Flask, request, jsonify
from flask_cors import CORS
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

app = Flask(__name__)
CORS(app)

# تنظیم مدل (مثال با مدل کوچک)
MODEL_NAME = "gpt2"  # برای تست - در تولید از مدل‌های فارسی استفاده کنید

class ChatBotAI:
    def __init__(self):
        self.tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
        self.model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
        self.history = []

    def generate_response(self, prompt, max_length=100):
        inputs = self.tokenizer.encode(prompt, return_tensors="pt")
        outputs = self.model.generate(
            inputs, 
            max_length=max_length,
            num_return_sequences=1,
            no_repeat_ngram_size=2,
            temperature=0.7,
            top_p=0.9
        )
        response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
        return response

chatbot = ChatBotAI()

@app.route('/api/chat', methods=['POST'])
def chat():
    data = request.json
    message = data.get('message', '')
    
    if not message:
        return jsonify({"error": "پیامی دریافت نشد"}), 400
    
    # تولید پاسخ
    response = chatbot.generate_response(message)
    
    # حذف prompt از پاسخ
    response = response.replace(message, "").strip()
    
    return jsonify({
        "response": response,
        "status": "success"
    })

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

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000, debug=False)