import os
import json
import time
from threading import Lock, Timer
from pydantic import BaseModel, Field
from mq_http_sdk.mq_client import MQClient
from mq_http_sdk.mq_producer import TopicMessage

from awe.util.logging import logger
from awe.types.common import Payload, ResponseMessage
from awe.util.workflow import convert_to_old_data, convert_to_old_result

def _env_flag(value) -> bool:
  return str(value).lower() in {'1', 'true', 'yes', 'y', 'on'}

class MQProducerOptions(BaseModel):
  endpoint: str = Field(title="MQ Endpoint", description="MQ Endpoint", default="")
  access_key: str = Field(title="MQ Access Key", description="MQ Access Key", default="")
  access_secret: str = Field(title="MQ Access Secret", description="MQ Access Secret", default="")
  instance_id: str = Field(title="实例ID", description="实例ID", default="")
  is_static_topic: bool = Field(title="是否为静态Topic", description="是否为静态Topic", default=False)
  env: str = Field(title="环境", description="环境", default="")
  appname: str = Field(title="应用名", description="应用名", default="")
  response_topic: str = Field(title="响应Topic", description="响应Topic", default="")
  progress_throttle_seconds: float = Field(title="进度限频秒数", description="进度限频秒数", default=2.0)
  global_max_qps: float = Field(title="全局最大发送QPS", description="全局最大发送QPS", default=20.0)

class MQProducer:
  def __init__(self, options: MQProducerOptions | None = None):
    self.options = options or MQProducerOptions(
      endpoint=os.getenv('ROCKETMQ_ENDPOINT') or '',
      access_key=os.getenv('ROCKETMQ_ACCESS_KEY') or '',
      access_secret=os.getenv('ROCKETMQ_ACCESS_SECRET') or '',
      instance_id=os.getenv('ROCKETMQ_INSTANCE_ID') or '',
      is_static_topic=_env_flag(os.getenv('ROCKETMQ_STATIC_TOPIC')),
      env=os.getenv('ENV') or '',
      appname=os.getenv('APPNAME') or '',
      response_topic=os.getenv('ROCKETMQ_BASE_RESPONSE_TOPIC') or '',
      progress_throttle_seconds=float(os.getenv('PROGRESS_THROTTLE_SECONDS', '2.0')),
      global_max_qps=float(os.getenv('GLOBAL_MAX_QPS', '20.0'))
    )
    self._progress_lock = Lock()
    self._progress_publish_at = {}
    self._trailing_timers = {}
    self._trailing_messages = {}
    self._global_lock = Lock()
    self._global_publish_at = None
    self.producer = self.get_producer(self.options)
    
  def get_producer(self, options: MQProducerOptions):
    mq_client = MQClient(
      options.endpoint,
      options.access_key,
      options.access_secret
    )
    if options.is_static_topic:
      producer_topic = options.response_topic
    else:
      topic_prefix = f"{options.env.upper()}_{options.appname.upper()}"
      producer_topic = f"{topic_prefix}_{options.response_topic}"
    
    logger.info(f'producer_topic is {producer_topic}')
    return mq_client.get_producer(options.instance_id, producer_topic)

  def send_progress_message(self, message: ResponseMessage, payload: Payload):
    message_body = {
      "code": "0",
      "status": "PROCESSING",
      "success": True,
      "taskId": payload.taskId,
      "userId": payload.userId,
      "clientId": payload.clientId,
      "client_id": payload.clientId,
      "appName": self.options.appname,
      "response": message.response_dict(),
      "workflowCode": payload.workflowCode
    }
    # 兼容旧版本，response_message 中要包含 progress
    if payload.messageTag != 'WEB':
      message_body["progress"] = message.response_dict().get("progress")
    # 所有进度消息统一走节流，避免高频发送触发 MQ 限流
    self._throttled_progress_publish(message, payload, json.dumps(message_body))
    
  def _progress_key(self, payload: Payload, message: ResponseMessage):
    response = message.response_dict()
    progress_data = response.get('progressData', {}) if isinstance(response, dict) else {}
    message_type = progress_data.get('type') or message.type or 'progress'
    return (payload.taskId, payload.messageTag, message_type)

  def _should_publish_progress(self, payload: Payload, message: ResponseMessage):
    throttle_seconds = max(float(self.options.progress_throttle_seconds or 0), 0)
    if throttle_seconds == 0:
      return True

    key = self._progress_key(payload, message)
    now = time.monotonic()
    with self._progress_lock:
      last_publish_at = self._progress_publish_at.get(key)
      if last_publish_at is not None and now - last_publish_at < throttle_seconds:
        return False
      self._progress_publish_at[key] = now
      return True

  def _schedule_trailing_publish(self, key, tag: str, delay: float):
    """在节流窗口结束后，自动发送最后一条被抑制的进度消息，确保最新进度不丢失"""
    def _flush():
      pending_message = None
      with self._progress_lock:
        self._trailing_timers.pop(key, None)
        pending_message = self._trailing_messages.pop(key, None)
        if pending_message is not None:
          self._progress_publish_at[key] = time.monotonic()
      if pending_message is not None:
        self.publish_message(pending_message, tag)

    old_timer = self._trailing_timers.pop(key, None)
    if old_timer is not None:
      old_timer.cancel()
    timer = Timer(delay, _flush)
    timer.daemon = True
    self._trailing_timers[key] = timer
    timer.start()

  def _clear_progress_state(self, payload: Payload):
    task_id = getattr(payload, 'taskId', None)
    if task_id is None:
      return
    with self._progress_lock:
      keys_to_remove = [key for key in self._progress_publish_at if key[0] == task_id]
      for key in keys_to_remove:
        self._progress_publish_at.pop(key, None)
        self._trailing_messages.pop(key, None)
        timer = self._trailing_timers.pop(key, None)
        if timer is not None:
          timer.cancel()

  def _throttled_progress_publish(self, message: ResponseMessage, payload: Payload, message_body: str):
    """节流 + trailing：被抑制的最后一条消息会在窗口结束后自动补发"""
    throttle_seconds = max(float(self.options.progress_throttle_seconds or 0), 0)
    if throttle_seconds == 0:
      self.publish_message(message_body, payload.messageTag)
      return

    key = self._progress_key(payload, message)
    now = time.monotonic()
    should_publish = False
    remaining_delay = 0.0

    with self._progress_lock:
      last_publish_at = self._progress_publish_at.get(key)
      if last_publish_at is None or now - last_publish_at >= throttle_seconds:
        self._progress_publish_at[key] = now
        self._trailing_messages.pop(key, None)
        should_publish = True
      else:
        # 被节流：暂存消息，等窗口结束后 trailing 发送
        self._trailing_messages[key] = message_body
        remaining_delay = throttle_seconds - (now - last_publish_at)

    if should_publish:
      # 取消该 key 上已有的 trailing timer
      with self._progress_lock:
        old_timer = self._trailing_timers.pop(key, None)
      if old_timer is not None:
        old_timer.cancel()
      self.publish_message(message_body, payload.messageTag)
    else:
      logger.debug(f"Throttled progress message, task_id: {payload.taskId}, trailing in {remaining_delay:.2f}s")
      self._schedule_trailing_publish(key, payload.messageTag, remaining_delay)

  def debounce_progress_message(self, message: ResponseMessage, payload: Payload, message_body: str):
    """保留向后兼容，内部统一走 _throttled_progress_publish"""
    self._throttled_progress_publish(message, payload, message_body)

  def send_error_message(self, message: ResponseMessage, payload: Payload):
    self._clear_progress_state(payload)
    message_body = {
      "code": message.code if message.code else "500",
      "status": "FAILED",
      "success": False,
      "taskId": payload.taskId,
      "userId": payload.userId,
      "clientId": payload.clientId,
      "client_id": payload.clientId,
      "appName": self.options.appname,
      "response": message.response_dict(),
      "workflowCode": payload.workflowCode
    }
    self.publish_message(json.dumps(message_body), tag=payload.messageTag)

  def send_complete_message(self, message: ResponseMessage, payload: Payload):
    self._clear_progress_state(payload)
    code = f"{message.code}" if message.code else "0"
    status = "SUCCEED" if code == "0" else "FAILED"
    message_body = {
      "code": code,
      "status": status,
      "success": status == "SUCCEED",
      "taskId": payload.taskId,
      "userId": payload.userId,
      "clientId": payload.clientId,
      "client_id": payload.clientId,
      "appName": self.options.appname,
      "workflowCode": payload.workflowCode,
      "response": message.response_dict(),
    }
    if payload.messageTag == 'WEB':
      message_body['results'] = []
    # 兼容老版本，results 是 response 的 data 的 images 的 output
    try:
      message_body['results'] = convert_to_old_result(message)
    except Exception as e:
      message_body['results'] = []
      logger.error(f'[MQProducer] for compatability, get image output from response error: {e}')
    # 兼容老版本，message_body 中包含 data
    try:
      message_body['data'] = convert_to_old_data(message)
    except Exception as e:
      message_body['data'] = {}
      logger.error(f'[MQProducer] for compatability, get image output from response error: {e}')
    self.publish_message(json.dumps(message_body), tag=payload.messageTag)
  
  def _check_global_rate(self) -> bool:
    """全局 QPS 限制，跨所有 task 的兜底保护"""
    max_qps = max(float(self.options.global_max_qps or 0), 0)
    if max_qps == 0:
      return True
    min_interval = 1.0 / max_qps
    now = time.monotonic()
    with self._global_lock:
      if self._global_publish_at is not None and now - self._global_publish_at < min_interval:
        return False
      self._global_publish_at = now
      return True

  def publish_message(self, message: str, tag: str | None = None):
    if not self._check_global_rate():
      logger.warning("Global QPS limit reached, dropping message to protect MQ")
      return
    try:
      topic_message = TopicMessage(message, tag or "")
      re_msg = self.producer.publish_message(topic_message)
      logger.info("Publish Succeed, TopicName:%s MessageID:%s MessageBody:%s" % (self.options.response_topic, re_msg.message_id, message))
    except Exception as e:
      logger.error("Publish, Message Fail. Exception:%s" % e)
