import os
import sys
from flask import Flask, request, jsonify
import torch
from transformers import BertTokenizer, BertForSequenceClassification
import joblib
import re
from functools import lru_cache
from concurrent.futures import ThreadPoolExecutor
from typing import List, Dict
import threading
import logging
import atexit
from datetime import datetime
import time

# ==================== 从环境变量读取配置 ====================
SERVICE_NAME = os.environ.get('SERVICE_NAME', 'address_classifier')
SERVICE_PORT = int(os.environ.get('SERVICE_PORT', 5003))
MODEL_ID = os.environ.get('MODEL_ID', 'circles1/address_model')
MODEL_DIR = os.environ.get('MODEL_DIR', '/app/services/bert_address/model')

# 模型文件路径
LABEL_ENCODER_PATH = os.path.join(MODEL_DIR, 'label_encoder.pkl')
# 如果上面路径不存在，尝试其他可能的文件名
if not os.path.exists(LABEL_ENCODER_PATH):
    import glob

    encoder_files = glob.glob(os.path.join(MODEL_DIR, 'label_encoder*.pkl'))
    if encoder_files:
        LABEL_ENCODER_PATH = encoder_files[0]

# ==================== 配置参数 ====================
MAX_LENGTH = 512
DEVICE = os.environ.get('DEVICE', 'cuda' if torch.cuda.is_available() else 'cpu')
# 如果设备设置为cuda但不可用，回退到cpu
if DEVICE == 'cuda' and not torch.cuda.is_available():
    DEVICE = 'cpu'

MAX_WORKERS = int(os.environ.get('MAX_WORKERS', 4))
CACHE_SIZE = int(os.environ.get('CACHE_SIZE', 2000))
TOKEN_CACHE_SIZE = int(os.environ.get('TOKEN_CACHE_SIZE', 1000))

# ==================== 日志配置 ====================
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# 打印启动信息
logger.info("=" * 50)
logger.info(f"Starting service: {SERVICE_NAME}")
logger.info(f"Service port: {SERVICE_PORT}")
logger.info(f"Model ID: {MODEL_ID}")
logger.info(f"Model directory: {MODEL_DIR}")
logger.info(f"Label encoder: {LABEL_ENCODER_PATH}")
logger.info(f"Device: {DEVICE}")
logger.info(f"Max workers: {MAX_WORKERS}")
logger.info("=" * 50)

app = Flask(__name__)
app.config['JSON_AS_ASCII'] = False

# 全局变量锁
model_lock = threading.Lock()


class Predictor:
    """预测器类，用于管理模型和分词器的生命周期"""
    _instance = None

    def __new__(cls):
        if cls._instance is None:
            with model_lock:
                if cls._instance is None:
                    cls._instance = super().__new__(cls)
                    cls._instance._initialized = False
        return cls._instance

    def __init__(self):
        if self._initialized:
            return

        logger.info(f"Initializing Predictor on device: {DEVICE}")
        start_time = time.time()

        try:
            # 检查模型目录是否存在
            if not os.path.exists(MODEL_DIR):
                raise FileNotFoundError(f"Model directory not found: {MODEL_DIR}")

            # 检查label_encoder是否存在
            if not os.path.exists(LABEL_ENCODER_PATH):
                raise FileNotFoundError(f"Label encoder not found: {LABEL_ENCODER_PATH}")

            # 加载模型和分词器
            logger.info(f"Loading tokenizer and model from {MODEL_DIR}...")
            self.tokenizer = BertTokenizer.from_pretrained(MODEL_DIR)
            self.model = BertForSequenceClassification.from_pretrained(MODEL_DIR).to(DEVICE)
            self.model.eval()
            logger.info("Model loaded successfully")

            # 加载标签映射器
            logger.info(f"Loading label encoder from {LABEL_ENCODER_PATH}...")
            self.label_encoder = joblib.load(LABEL_ENCODER_PATH)
            self.label_map = {i: label for i, label in enumerate(self.label_encoder.classes_)}
            logger.info(f"Loaded {len(self.label_map)} labels: {list(self.label_map.values())}")

            # 创建线程池
            self.executor = ThreadPoolExecutor(max_workers=MAX_WORKERS)
            logger.info(f"ThreadPoolExecutor created with {MAX_WORKERS} workers")

            # 注册关闭钩子
            atexit.register(self._shutdown)

            # 模型预热
            self._warmup()

            self._initialized = True
            elapsed_time = time.time() - start_time
            logger.info(f"Predictor initialized successfully in {elapsed_time:.2f}s")

        except Exception as e:
            logger.error(f"Failed to initialize Predictor: {e}")
            raise

    def _shutdown(self):
        """清理资源"""
        logger.info("Shutting down executor...")
        self.executor.shutdown(wait=True)
        logger.info("Executor shutdown complete")

    def _warmup(self):
        """模型预热，确保CUDA初始化"""
        logger.info("Warming up model...")
        try:
            test_text = "测试地址"
            self.predict_single(test_text)
            logger.info("Model warmup completed")
        except Exception as e:
            logger.warning(f"Model warmup failed: {e}")

    @lru_cache(maxsize=CACHE_SIZE)
    def clean_text(self, text: str) -> str:
        """文本清洗函数，带缓存"""
        if not isinstance(text, str):
            return ""
        # 只保留中文字符
        cleaned_text = re.sub(r'[^\u4e00-\u9fa5]', '', text)
        return cleaned_text.strip()

    def predict_single(self, text: str) -> Dict:
        """单个文本预测"""
        if not text or not isinstance(text, str):
            return {"address": "", "confidence": 0.0, "error": "Invalid input"}

        try:
            cleaned_text = self.clean_text(text)
            if not cleaned_text:
                logger.warning(f"Empty text after cleaning: {text[:50]}...")
                return {"address": "", "confidence": 0.0}

            inputs = self.tokenizer(
                cleaned_text,
                return_tensors="pt",
                truncation=True,
                padding=True,
                max_length=MAX_LENGTH
            ).to(DEVICE)

            with torch.no_grad():
                outputs = self.model(**inputs)
                probs = torch.softmax(outputs.logits, dim=1).cpu()
                top_prob, top_idx = torch.topk(probs, k=1)

            return {
                "address": self.label_map[top_idx.item()],
                "confidence": round(top_prob.item(), 4)
            }

        except torch.cuda.OutOfMemoryError as e:
            logger.error(f"CUDA OOM error: {e}")
            torch.cuda.empty_cache()
            return {"address": "", "confidence": 0.0, "error": "GPU memory exhausted"}
        except Exception as e:
            logger.error(f"Prediction error for text '{text[:50]}...': {e}")
            return {"address": "", "confidence": 0.0, "error": str(e)}

    def batch_predict(self, texts: List[str]) -> List[Dict]:
        """批量预测（并发处理）"""
        if not texts:
            return []

        logger.info(f"Processing {len(texts)} texts")
        start_time = time.time()

        # 使用线程池并发处理
        futures = [self.executor.submit(self.predict_single, text) for text in texts]
        results = [future.result() for future in futures]

        elapsed_time = time.time() - start_time
        logger.info(f"Batch prediction completed in {elapsed_time:.2f}s for {len(texts)} texts")

        return results

    def get_stats(self) -> Dict:
        """获取预测器统计信息"""
        return {
            "service_name": SERVICE_NAME,
            "model_id": MODEL_ID,
            "model_dir": MODEL_DIR,
            "device": DEVICE,
            "num_labels": len(self.label_map),
            "labels": list(self.label_map.values()),
            "max_length": MAX_LENGTH,
            "max_workers": MAX_WORKERS
        }


# 初始化预测器
try:
    predictor = Predictor()
    logger.info("Predictor instance created successfully")
except Exception as e:
    logger.error(f"Failed to create predictor instance: {e}")
    raise


@app.route('/predict', methods=['POST'])
def predict():
    """预测接口"""
    try:
        data = request.get_json()
        if not data:
            return jsonify({
                "status": "error",
                "error": "Invalid request, JSON body required"
            }), 400

        if 'address' not in data:
            return jsonify({
                "status": "error",
                "error": "Invalid request, 'address' is required"
            }), 400

        addresses = data['address']
        if not isinstance(addresses, list):
            addresses = [addresses]

        # 限制最大处理数量
        max_items = data.get('max_items', 1000)
        if len(addresses) > max_items:
            logger.warning(f"Too many items: {len(addresses)} > {max_items}, truncating")
            addresses = addresses[:max_items]

        # 去重（可选）
        if data.get('deduplicate', False):
            addresses = list(dict.fromkeys(addresses))
            logger.info(f"After deduplication: {len(addresses)} unique texts")

        # 批量预测
        results = predictor.batch_predict(addresses)

        return jsonify({
            "status": "success",
            "predictions": results,
            "metadata": {
                "total": len(results),
                "timestamp": datetime.now().isoformat(),
                "service": SERVICE_NAME,
                "device": DEVICE
            }
        })

    except Exception as e:
        logger.error(f"Prediction endpoint error: {e}", exc_info=True)
        return jsonify({
            "status": "error",
            "error": f"Internal server error: {str(e)}"
        }), 500


@app.route('/health', methods=['GET'])
def health_check():
    """健康检查接口"""
    try:
        stats = predictor.get_stats()
        return jsonify({
            "status": "healthy",
            "timestamp": datetime.now().isoformat(),
            "stats": stats
        })
    except Exception as e:
        return jsonify({
            "status": "unhealthy",
            "error": str(e)
        }), 500


@app.route('/stats', methods=['GET'])
def get_stats():
    """获取预测器统计信息"""
    try:
        stats = predictor.get_stats()
        return jsonify({
            "status": "success",
            "stats": stats
        })
    except Exception as e:
        return jsonify({
            "status": "error",
            "error": str(e)
        }), 500


@app.errorhandler(404)
def not_found(error):
    return jsonify({
        "status": "error",
        "error": "Endpoint not found"
    }), 404


@app.errorhandler(500)
def internal_error(error):
    return jsonify({
        "status": "error",
        "error": "Internal server error"
    }), 500


if __name__ == '__main__':
    logger.info(f"Starting Flask application: {SERVICE_NAME}")
    logger.info(f"Model path: {MODEL_DIR}")
    logger.info(f"Device: {DEVICE}")
    logger.info(f"Port: {SERVICE_PORT}")

    app.run(
        host='0.0.0.0',
        port=SERVICE_PORT,
        threaded=True,
        debug=False
    )

else:
    # 兼容 WSGI 标准（如 Gunicorn）
    application = app
    logger.info(f"Application loaded for WSGI server: {SERVICE_NAME}")