"""
VLM图像处理工具模块
包含图像识别、图像压缩、图像保存等功能
"""
import os
import re
import json
import math
import base64
import uuid
from datetime import datetime
from typing import Optional, Dict, Any
from PIL import Image
from io import BytesIO
import tiktoken
from langchain_core.tools import tool

from config import VLM_CONFIG, SERVER_CONFIG
from modelsAPI.model_api import OpenaiAPI
from utils.function_tracker import track_function_calls


VLM_API_BASE = VLM_CONFIG["api_base"]
VLM_MODEL_NAME = VLM_CONFIG["model"]
MAX_MODEL_LEN = 32768
INPUT_TOKEN_SAFETY_RATIO = 0.8
MIN_IMAGE_SIZE = 224
JPEG_QUALITY = 85

SYSTEM_PROMPT_VLM = """你是一个专业的工业设备巡检专家助手。
你的任务是协助工程师分析现场情况。

当你遇到以下情况时：
1. 用户提供了图片路径（如 .png, .jpg, /app/data/... 等）。
2. 用户询问有关设备外观、颜色、仪表读数或物理损坏的问题。

你应该：
- 调用 `vlm_image_to_text` 工具。
- 在 `prompt` 参数中，根据用户的具体意图（如：检查漏油、读仪表、看指示灯）量身定制分析指令。
- 获取 VLM 结果后，结合专业知识给出建议。
"""


def encode_image(image_path: str) -> str:
    """将本地图片转换为 Base64 编码"""
    with open(image_path, "rb") as f:
        return base64.b64encode(f.read()).decode("utf-8")


def estimate_qwen_vl_visual_tokens(width: int, height: int) -> int:
    """
    根据 Qwen-VL 的图像分块策略估算视觉 tokens。
    公式来源：Qwen-VL 官方技术文档（Naive resizing + 192 tokens per unit）
    """
    units_h = math.ceil(height / 336)
    units_w = math.ceil(width / 336)
    total_units = units_h * units_w + 1
    return total_units * 192


def count_text_tokens(text: str) -> int:
    """使用 tiktoken 估算文本 token 数（兼容中英文）"""
    enc = tiktoken.get_encoding("cl100k_base")
    return len(enc.encode(text))


def smart_resize_base64_image(
    base64_image: str,
    prompt_text: str,
    max_input_tokens: int = int(MAX_MODEL_LEN * INPUT_TOKEN_SAFETY_RATIO),
    min_size: int = MIN_IMAGE_SIZE,
    jpeg_quality: int = JPEG_QUALITY
) -> str:
    """
    智能压缩 base64 图像，确保 (文本 tokens + 视觉 tokens) <= max_input_tokens
    """
    if not base64_image.startswith("data:image"):
        raise ValueError("Invalid base64 image format. Must start with 'data:image'.")

    try:
        header, b64_data = base64_image.split(",", 1)
        image_data = base64.b64decode(b64_data)
        img = Image.open(BytesIO(image_data)).convert("RGB")
        orig_w, orig_h = img.size
    except Exception as e:
        raise ValueError(f"Failed to decode base64 image: {e}")

    text_tokens = count_text_tokens(prompt_text)
    visual_tokens = estimate_qwen_vl_visual_tokens(orig_w, orig_h)
    total_tokens = text_tokens + visual_tokens

    print(f"[VLM Preprocess] 原始尺寸: {orig_w}x{orig_h} | 文本 tokens: {text_tokens} | 视觉 tokens: {visual_tokens} | 总计: {total_tokens}")

    if total_tokens <= max_input_tokens:
        print("[VLM Preprocess] ✅ 无需压缩")
        return base64_image

    scale = 1.0
    target_tokens = total_tokens
    while target_tokens > max_input_tokens and scale > 0.1:
        scale *= 0.8
        new_w = max(min_size, int(orig_w * scale))
        new_h = max(min_size, int(orig_h * scale))
        target_tokens = text_tokens + estimate_qwen_vl_visual_tokens(new_w, new_h)

    final_w = max(min_size, int(orig_w * scale))
    final_h = max(min_size, int(orig_h * scale))

    print(f"[VLM Preprocess] ⚠️ 需压缩 → 新尺寸: {final_w}x{final_h} | 预估总 tokens: {target_tokens}")

    resized_img = img.resize((final_w, final_h), Image.Resampling.LANCZOS)
    buffer = BytesIO()
    resized_img.save(buffer, format="JPEG", quality=jpeg_quality)
    compressed_b64 = base64.b64encode(buffer.getvalue()).decode("utf-8")
    return f"data:image/jpeg;base64,{compressed_b64}"


def save_base64_image_and_record(
    base64_image: str,
    description: str,
    output_dir: str = "image_input",
    record_file: str = "image_input/image_records.json"
) -> Dict[str, Any]:
    """
    保存 base64 图片到指定目录，并记录图片名称和描述到 JSON 文件。
    
    Args:
        base64_image: base64 编码的图片数据（data:image/xxx;base64,... 格式）
        description: 图片的描述文本
        output_dir: 图片保存目录
        record_file: 记录 JSON 文件路径
    
    Returns:
        包含保存结果的字典：{"success": bool, "image_name": str, "image_path": str, "error": str}
    """
    try:
        if not base64_image.startswith("data:image"):
            return {"success": False, "error": "Invalid base64 image format"}
        
        header, b64_data = base64_image.split(",", 1)
        
        if "jpeg" in header or "jpg" in header:
            ext = "jpg"
        elif "png" in header:
            ext = "png"
        elif "gif" in header:
            ext = "gif"
        elif "webp" in header:
            ext = "webp"
        else:
            ext = "jpg"
        
        os.makedirs(output_dir, exist_ok=True)
        
        image_name = f"{uuid.uuid4().hex[:12]}_{datetime.now().strftime('%Y%m%d%H%M%S')}.{ext}"
        image_path = os.path.join(output_dir, image_name)
        
        image_data = base64.b64decode(b64_data)
        with open(image_path, "wb") as f:
            f.write(image_data)
        
        record_entry = {
            "image_name": image_name,
            "image_path": image_path,
            "description": description,
            "timestamp": datetime.now().isoformat()
        }
        
        os.makedirs(os.path.dirname(record_file), exist_ok=True)
        
        existing_records = []
        if os.path.exists(record_file):
            try:
                with open(record_file, "r", encoding="utf-8") as f:
                    existing_records = json.load(f)
                    if not isinstance(existing_records, list):
                        existing_records = []
            except (json.JSONDecodeError, IOError):
                existing_records = []
        
        existing_records.append(record_entry)
        
        with open(record_file, "w", encoding="utf-8") as f:
            json.dump(existing_records, f, ensure_ascii=False, indent=2)
        
        print(f"[ImageSaved] 图片已保存: {image_path}")
        
        return {
            "success": True,
            "image_name": image_name,
            "image_path": image_path,
            "error": None
        }
        
    except Exception as e:
        return {
            "success": False,
            "image_name": None,
            "image_path": None,
            "error": str(e)
        }


@tool
async def vlm_image_to_text(
    image_path: str,
    prompt: Optional[str] = None
) -> str:
    """
    调用 Qwen3.5-35B 进行工业巡检分析。
    - 自动移除所有 <think...</think 包裹的内容（包括多行）
    - 仅返回最终结果文本
    - 支持 40K token 上下文，自动压缩大图
    """

    default_prompt = (
        "【重要指令】禁止输出任何形式的   <think...</think、<reasoning></reasoning> 或其他内部思考过程。"
        "仅输出最终的专业分析结果。\n\n"
        "你是一位资深的工业巡检专家。请仔细观察这张图片并按以下结构报告：\n"
        "1. 设备识别：图中是什么设备？型号或标签是否可见？\n"
        "2. 状态评估：仪表盘数值是否在正常区间？指示灯颜色是什么？\n"
        "3. 异常检测：是否存在漏油、生锈、电缆破损、螺丝松动或烟雾？\n"
        "4. 环境风险：周围是否有杂物堆放或安全隐患？\n"
        "请用专业、客观、简短的语言描述。"
    )

    final_prompt = f"{default_prompt}\n特别注意：{prompt}" if prompt else default_prompt

    safe_image = smart_resize_base64_image(
        base64_image=image_path,
        prompt_text=final_prompt,
        max_input_tokens=int(MAX_MODEL_LEN * INPUT_TOKEN_SAFETY_RATIO)
    )

    try:
        result = await OpenaiAPI.open_api_vl_without_thinking(safe_image)

        print(f"✅ 最终输出:\n{result}")
        
        save_result = save_base64_image_and_record(
            base64_image=image_path,
            description=result,
            output_dir="picture",
            record_file="picture/image_records.json"
        )
        if save_result["success"]:
            print(f"📷 图片已保存: {save_result['image_name']}")
            image_url = f"{SERVER_CONFIG['base_url']}/picture/{save_result['image_name']}"
            result = f"{result}\n\n![用户上传图片]({image_url})"
        else:
            print(f"⚠️ 图片保存失败: {save_result['error']}")
        
        return result

    except Exception as e:
        import traceback
        error_msg = f"VLM 调用失败: {str(e)}\n{traceback.format_exc()}"
        print(error_msg)
        return error_msg


@track_function_calls
async def image_rag_describe(
    image_path: str,
    query: str,
    top_k: int = 5
) -> str:
    """
    图像RAG描述工具
    先用VLM分析图像，然后用RAG检索相关知识

    Args:
        image_path: 图像路径（base64格式）
        query: 查询文本
        top_k: RAG检索数量

    Returns:
        分析结果文本
    """
    try:
        vlm_result = await vlm_image_to_text(image_path, prompt=query)
        
        from .rag_tools import rag_search
        rag_result = await rag_search(query=f"{query} {vlm_result}", top_k=top_k)
        
        if rag_result.get("success", False):
            source_citation = rag_result.get("sourceCitation", {})
            rag_text = ""
            for resource, items in source_citation.items():
                for item in items:
                    rag_text += f"【{resource}】{item.get('text', '')}\n\n"
            
            return f"图像分析结果：\n{vlm_result}\n\n相关知识：\n{rag_text}"
        else:
            return vlm_result
            
    except Exception as e:
        return f"图像RAG分析失败: {str(e)}"