import json
import time
import re
import os
import requests
from consts.consts import Qwen
from config import config
from utils.logging import logging
from utils.local import get_local_metrics
from models.user_preference import UserPreference
from errors.errors import AlgorithmException, BadRequestException, UnauthorizedException
from consts.consts import GenSQLErrorMsg
from tenacity import retry, stop_after_attempt, wait_random_exponential, before_sleep_log

# Import DashScope SDK (for Qwen API calls)
# If not installed, will be mocked in tests
try:
    import dashscope
    from dashscope import Generation
except ImportError:
    dashscope = None
    Generation = None

logger = logging.getLogger(__name__)

# Billing2 service constants
BILLING_SOURCE_HEADER_KEY = "X-BILLING-REQUEST-SOURCE"
BILLING_SOURCE_OPS = "ops"

# Qwen model mapping
qwen_model_id = {
    Qwen.DEFAULT: "qwen-plus",  # Default to qwen-plus (128K context)
    Qwen.TURBO: "qwen-turbo",
    Qwen.PLUS: "qwen-plus",
    Qwen.MAX: "qwen-max"
}


def invoke_qwen_with_text(prompt, output_format, message_list=[], dry_run=False, ai_model=Qwen.DEFAULT, extra_info={}):
    """
    Invoke Alibaba Cloud Qwen model for inference

    Args:
        prompt: Prompt text
        output_format: Output format ("json" or other)
        message_list: Message history list (for multi-turn conversations)
        dry_run: Dry run mode (only check token limits)
        ai_model: Qwen model ID (qwen-turbo/qwen-plus/qwen-max)
        extra_info: Extra information (module, session_id, query_id, etc.)

    Returns:
        Parsed JSON object (if output_format="json") or text string

    Alibaba Cloud DashScope documentation:
    https://help.aliyun.com/zh/dashscope/developer-reference/api-details
    """
    model_id = qwen_model_id.get(ai_model, "").strip()
    if not model_id:
        logging.error(
            "session_id = %s, query_id = %s, module = %s, %s is not supported",
            extra_info.get("session_id"),
            extra_info.get("query_id"),
            extra_info.get("module"),
            ai_model
        )
        raise BadRequestException(f"Qwen model not supported: {ai_model}")

    session_id = extra_info.get("session_id")
    query_id = extra_info.get("query_id")
    module = extra_info.get("module")

    time_before = int(time.time())
    logging.info(
        "session_id = %s, query_id = %s, module = %s, start Qwen request, "
        "model_id = %s, prompt_length = %s, dry_run = %s",
        session_id, query_id, module, model_id, len(prompt), str(dry_run)
    )

    # Get API Key (two-tier configuration: user-level > system-level)
    api_key = _get_qwen_api_key(extra_info)

    # Dry run: only check token limits
    if dry_run:
        # TODO: implement token count check
        return True

    # Get org_id for building euid header
    metrics = get_local_metrics()
    org_id = extra_info.get("org_id") or metrics.get("org_id")

    # Call Qwen API
    try:
        resp_str, input_tokens, output_tokens = _call_qwen_api(
            model_id, prompt, message_list, api_key, org_id
        )

        time_cost = int(time.time()) - time_before
        logging.info(
            "session_id = %s, query_id = %s, module = %s, end Qwen request, "
            "model_id = %s, response_length = %s, input_tokens = %s, "
            "output_tokens = %s, time_cost = %ss",
            session_id, query_id, module, model_id, len(resp_str),
            input_tokens, output_tokens, time_cost
        )

        # Extract and fix JSON (reuse claude's logic)
        from svc.claude import extract_text
        resp_json = extract_text(resp_str)
        logging.info(
            "session_id = %s, query_id = %s, module = %s, extract result finished, resp = %s",
            session_id, query_id, module, resp_json
        )

        # Build message history
        messages = message_list + [{"role": "user", "content": prompt}]
        messages.append({"role": "assistant", "content": resp_json})

        # Handle JSON format output
        if output_format == "json":
            try:
                return json.loads(resp_json)
            except json.JSONDecodeError as e:
                logging.warning(
                    "session_id = %s, query_id = %s, module = %s, "
                    "JSON decode failed, attempting repair, error = %s",
                    session_id, query_id, module, str(e)
                )
                # Try to repair JSON
                from json_repair import repair_json
                try:
                    repaired = repair_json(resp_json)
                    return json.loads(repaired)
                except Exception as repair_error:
                    logging.error(
                        "session_id = %s, query_id = %s, module = %s, "
                        "JSON repair failed, error = %s",
                        session_id, query_id, module, str(repair_error)
                    )
                    # Recursively call itself to try regeneration
                    return invoke_qwen_with_text(
                        f"Your previous response was not valid JSON. "
                        f"Please provide a valid JSON response for: {prompt}",
                        output_format, messages, dry_run, ai_model, extra_info
                    )
        else:
            return resp_json

    except UnauthorizedException:
        raise
    except BadRequestException:
        raise
    except Exception as e:
        logging.error(
            "session_id = %s, query_id = %s, module = %s, Qwen API error, error = %s",
            session_id, query_id, module, str(e)
        )
        raise AlgorithmException(f"Qwen API error: {str(e)}")


def _get_qwen_api_key(extra_info):
    """
    Get Qwen API Key
    Priority: user-defined > system environment variable

    Args:
        extra_info: Extra information containing app_id

    Returns:
        API Key string

    Raises:
        BadRequestException: API Key not configured
    """
    # System-level default
    api_key = config.QWEN_API_KEY
    key_source = "system"

    # Check user-defined API Key (higher priority)
    metrics = get_local_metrics()
    app_id = metrics.get("app_id", 0)
    if app_id:
        preference = UserPreference.load(app_id)
        if preference.llm_api_key:
            api_key = preference.llm_api_key
            key_source = "user"
            logging.info(
                "session_id = %s, query_id = %s, module = %s, "
                "using custom Qwen API key for app_id = %s",
                extra_info.get("session_id"),
                extra_info.get("query_id"),
                extra_info.get("module"),
                app_id
            )

    if not api_key:
        logging.error(
            "session_id = %s, query_id = %s, module = %s, Qwen API key not configured",
            extra_info.get("session_id"),
            extra_info.get("query_id"),
            extra_info.get("module")
        )
        raise BadRequestException(
            "Qwen API key not configured. Please set QWEN_API_KEY environment "
            "variable or provide custom key in user preferences."
        )

    logging.info(
        "session_id = %s, query_id = %s, module = %s, Qwen API key loaded, source = %s",
        extra_info.get("session_id"),
        extra_info.get("query_id"),
        extra_info.get("module"),
        key_source
    )

    return api_key


@retry(
    reraise=True,
    wait=wait_random_exponential(min=60, max=120),
    stop=stop_after_attempt(3),
    before_sleep=before_sleep_log(logger, logging.INFO)
)
def _call_qwen_api(model_id, prompt, message_list, api_key, org_id=None):
    """
    Call Qwen API (with retry mechanism)

    Args:
        model_id: Qwen model ID
        prompt: Prompt text
        message_list: Message history list
        api_key: API Key
        org_id: Organization ID (for building euid header)

    Returns:
        tuple: (response_text, input_tokens, output_tokens)

    Raises:
        UnauthorizedException: Authentication failed
        BadRequestException: Token limit exceeded
        AlgorithmException: Other errors
    """
    try:
        # Check if dashscope is available
        if dashscope is None or Generation is None:
            raise AlgorithmException(
                "DashScope SDK not installed. Please run: pip install dashscope>=1.14.0"
            )

        # Configure international endpoint (if user is in international region)
        # Can be customized via DASHSCOPE_BASE_URL environment variable
        base_url = os.getenv('DASHSCOPE_BASE_URL', 'https://dashscope-intl.aliyuncs.com/api/v1')
        dashscope.base_http_api_url = base_url

        logging.info("Qwen API endpoint configured, base_url = %s, model_id = %s", base_url, model_id)

        # Build message list
        messages = message_list + [{"role": "user", "content": prompt}]

        # Build euid header for billing tracking
        euid_headers = _build_dashscope_euid_header(org_id)
        has_euid = bool(euid_headers)
        logging.info("Qwen API calling, model_id = %s, messages_count = %s, prompt_length = %s, org_id = %s, has_euid = %s, euid_headers = %s",
                     model_id, len(messages), len(prompt), org_id, has_euid, euid_headers)

        # Call Qwen API
        response = Generation.call(
            api_key=api_key,
            model=model_id,
            messages=messages,
            result_format='message',  # Return message format
            stream=False,  # Don't use streaming output
            temperature=0,  # Deterministic output
            top_p=1.0,
            max_tokens=4096,  # Maximum output token count
            headers=euid_headers
        )

        logging.info("Qwen API response received, status_code = %s", response.status_code)

        # Check response status
        if response.status_code != 200:
            error_msg = response.message if hasattr(response, 'message') else str(response)
            logging.error("Qwen API error, status_code = %s, error_msg = %s", response.status_code, error_msg)

            # Authentication error
            if response.status_code == 401 or "Unauthorized" in error_msg:
                logging.error("Qwen API authentication failed, status_code = %s, error_msg = %s",
                              response.status_code, error_msg)
                raise UnauthorizedException(f"Qwen API authentication failed: {error_msg}")

            # Token limit exceeded
            if "token" in error_msg.lower() and "limit" in error_msg.lower():
                logging.error("Qwen API token limit exceeded, error_msg = %s", error_msg)
                raise BadRequestException(GenSQLErrorMsg.TOKEN_EXCEED_ERROR)

            # Other errors
            logging.error("Qwen API unknown error, status_code = %s, error_msg = %s",
                          response.status_code, error_msg)
            raise AlgorithmException(f"Qwen API error (status={response.status_code}): {error_msg}")

        # Parse response
        result = response.output
        usage = response.usage

        # Extract text content
        if hasattr(result, 'choices') and len(result.choices) > 0:
            resp_text = result.choices[0].message.content
            logging.info("Qwen API response extracted, response_length = %s", len(resp_text))
        else:
            logging.error("Qwen API returned empty response")
            raise AlgorithmException("Qwen API returned empty response")

        # 提取 token 使用情况
        input_tokens = usage.input_tokens if hasattr(usage, 'input_tokens') else 0
        output_tokens = usage.output_tokens if hasattr(usage, 'output_tokens') else 0

        logging.info("Qwen API token usage, input_tokens = %s, output_tokens = %s, total_tokens = %s",
                     input_tokens, output_tokens, input_tokens + output_tokens)

        return resp_text, input_tokens, output_tokens

    except (UnauthorizedException, BadRequestException, AlgorithmException):
        raise
    except Exception as e:
        raise AlgorithmException(f"Qwen API call failed: {str(e)}")


def _get_aliyun_subscription_record(org_id):
    """
    Get Aliyun subscription record from billing2 service

    Args:
        org_id: Organization/tenant ID

    Returns:
        dict: {"ali_uid": "xxx", "biz_id": "xxx", "in_trial": false, "is_subscribed": false}
        or None if request fails
    """
    if not org_id:
        return None

    try:
        url = f"{config.BILLING2_ADDR}/api/v1/billing/tenants/{org_id}/aliyun_subscription_record"
        headers = {BILLING_SOURCE_HEADER_KEY: BILLING_SOURCE_OPS}
        resp = requests.get(url, headers=headers, timeout=5)

        if resp.status_code != 200:
            logging.warning(
                "Failed to get aliyun subscription record, org_id = %s, status_code = %s",
                org_id, resp.status_code
            )
            return None

        return resp.json()
    except Exception as e:
        logging.warning(
            "Error getting aliyun subscription record, org_id = %s, error = %s",
            org_id, str(e)
        )
        return None


def _build_dashscope_euid_header(org_id):
    """
    Build x-dashscope-euid header for Qwen API calls

    Args:
        org_id: Organization/tenant ID

    Returns:
        dict: Headers dict with x-dashscope-euid, or empty dict if user not subscribed
    """
    # Get aliyun subscription info
    subscription = _get_aliyun_subscription_record(org_id)

    # Only add euid header if user is subscribed
    if not subscription or not subscription.get("is_subscribed"):
        return {}

    ali_uid = subscription.get("ali_uid", "")
    if not ali_uid:
        return {}

    euid_data = {
        "bizType": "B2B",
        "moduleType": "Third-partyproducts",
        "moduleCode": config.QWEN_MODULE_CODE,
        "accountType": "Aliyun",
        "accountId": ali_uid,
    }

    return {"x-dashscope-euid": json.dumps(euid_data)}


def validate_qwen_api_key(api_key, ai_model):
    """
    验证 Qwen API Key 是否有效

    Args:
        api_key: Qwen API Key
        ai_model: 模型 ID

    Returns:
        bool: True 表示有效，False 表示无效
    """
    try:
        model_id = qwen_model_id.get(ai_model, Qwen.TURBO)
        test_prompt = "Hello, this is a test."

        _call_qwen_api(model_id, test_prompt, [], api_key, None)
        return True
    except Exception as e:
        logging.error("Qwen API key validation failed, error = %s", str(e))
        return False