import numpy as np
from PIL import Image
import os
import asyncio
import hashlib
import json
import aiohttp
import logging
import base64
import cv2
import uuid
from openai import AsyncOpenAI, BadRequestError
from Config.Config import (
    ALY_LLM_API_KEY, VL_MODEL_NAME, VL_MODEL_NAME_AD
)
from Apps.AiTeJiYiChong.Config.Setting import (
    SAFE_EXCLUDE_RATIO, FALLBACK_WIDTH, FALLBACK_HEIGHT,
    BOTTOM_SAFE_EXCLUDE_RATIO
)
from Apps.AiTeJiYiChong import Kit

logger = logging.getLogger(__name__)


class ReadImageKit:
    _client = AsyncOpenAI(
        api_key=ALY_LLM_API_KEY,
        base_url="https://dashscope.aliyuncs.com/compatible-mode/v1"
    )

    # 通用回退设备信息，仅在无法动态获取设备信息时使用
    _FALLBACK_DEVICE_INFO = {
        "displayWidth": FALLBACK_WIDTH,
        "displayHeight": FALLBACK_HEIGHT,
        "productName": "generic"
    }

    _prompt = (
        "仅输出JSON数组（不含任何说明文字），按从左到右、从上到下的顺序识别图片中由【绿色方框】标识的充电站区域。识别规则如下：\n"
        "1. 必须是图中用绿色实线方框圈出的区域。\n"
        "2. 每一个卡片区域必须同时具备以下所有要素，否则严禁识别：\n"
        "   - 场站名称 (station_name)；\n"
        "   - 距离信息 (distance, 例如 '2.82km' 或 '90m')，通常位于卡片右侧蓝色胶囊区域内；\n"
        "   - 金额/电费 (price，例如 '1.2500')，通常以红色字体显示；\n"
        "   - 充电枪信息 (piles，包含'快'或'慢'的类型、总枪数和空闲枪数，例如 '快 4/4')。\n"
        "3. 如果绿色方框内缺少上述任何一项要素，说明它不是真正的场站卡片，请直接跳过。\n"
        "\n"
        "JSON对象字段要求：\n"
        "1. b_use: 状态标识（1或0）。如果场站名称为灰色或带有“暂停使用”等标签，则为0，否则为1。\n"
        "2. station_name: 场站名称；\n"
        "3. price: 一度电的价格（数字，如 1.2500）；\n"
        "4. piles: 充电枪列表 [{type: '快', free: 4, total: 4}]；\n"
        "5. parking: 停车费用描述（通常在蓝色'P'图标后，例如 '免费停车三小时'）；\n"
        "6. distance: 距离信息字符串（例如 '2.82km' 或 '90m'）；\n"
        "7. bounds: {x1,y1,x2,y2} 区域像素坐标（0-1000）；\n"
        "8. bounds_norm: {left,top,right,bottom} 归一化坐标(0-1)；\n"
        "9. station_name_bounds: 场站名称文字区域坐标 {x1,y1,x2,y2}（0-1000）；\n"
        "10. station_name_bounds_norm: 场站名称文字归一化坐标(0-1)。\n"
        "\n"
        "重要约束：\n"
        "A. 严禁识别未被绿色方框圈出的区域。如顶部的“长春市”选择、顶部的搜索框、以及中间的“推荐站点”等标签。\n"
        "B. 真正的场站卡片在绿色方框内包含：场站名称、金额（红色文字）、距离（右侧蓝色背景内）、充电枪状态（绿色或蓝色徽章，格式为 闲x/x 或 x/x）。\n"
        "C. 严禁将顶部的功能图标（如我的订单、收藏站点等）误认为场站卡片。\n"
        "\n"
        "严格返回纯JSON格式。"
    )

    @staticmethod
    def _extract_json(text: str) -> str:
        if not text:
            return "[]"

        cleaned = text.strip()
        if "```" in cleaned:
            lines = []
            for line in cleaned.splitlines():
                if line.strip().startswith("```"):
                    continue
                lines.append(line)
            cleaned = "\n".join(lines).strip()

        decoder = json.JSONDecoder()

        pos = 0
        while pos < len(cleaned):
            idx_dict = cleaned.find("{", pos)
            idx_list = cleaned.find("[", pos)

            candidates = [i for i in (idx_dict, idx_list) if i != -1]
            if not candidates:
                break

            start = min(candidates)
            snippet = cleaned[start:]
            try:
                _, end = decoder.raw_decode(snippet)
                return snippet[:end]
            except json.JSONDecodeError:
                pos = start + 1
                continue

        return "[]"

    _prompt_detail = (
        "仅输出JSON对象（不含任何说明文字），识别充电站详情图片中的以下信息：\n"
        "1. station_name: 场站名称；\n"
        "2. address: 场站完整地址（通常在定位图标旁）。\n"
        "\n"
        "特别说明：\n"
        "- 严禁输出 Markdown 代码块标签，严格返回纯JSON对象。"
    )

    @staticmethod
    def _to_minutes(t_str: str) -> int:
        """HH:MM -> 分钟数"""
        if not t_str or ":" not in t_str:
            return 0
        try:
            h, m = map(int, t_str.split(":"))
            return h * 60 + m
        except:
            return 0

    @staticmethod
    def _fmt(t: int) -> str:
        """分钟数 -> HH:MM"""
        h = t // 60
        m = t % 60
        return f"{h:02d}:{m:02d}"

    @staticmethod
    def expand_schedule_to_24h(rows: list) -> list:
        """
        将时段列表规整为全天24个整点小时段
        """
        # 预处理：转换为分钟区间
        intervals = []
        for r in rows:
            s = ReadImageKit._to_minutes(r.get("start"))
            e = ReadImageKit._to_minutes(r.get("end"))
            if e <= s and e != 0: # 处理 00:00-00:00 这种可能
                if e == 0: e = 1440
                else: continue
            if e == 0 and s > 0: e = 1440
            
            s = max(0, s)
            e = min(1440, e)
            intervals.append({
                "s": s, "e": e,
                "price": r.get("price")
            })
        
        # 排序
        intervals.sort(key=lambda x: (x["s"], x["e"]))

        result = []
        for h in range(24):
            hs = h * 60
            he = (h + 1) * 60
            best_price = None
            
            # 找到覆盖当前小时段的价格
            for it in intervals:
                # 计算重叠
                overlap_s = max(hs, it["s"])
                overlap_e = min(he, it["e"])
                if overlap_e > overlap_s:
                    best_price = it["price"]
                    break # 简单处理，取第一个覆盖的
            
            result.append({
                "start": ReadImageKit._fmt(hs),
                "end": ReadImageKit._fmt(he),
                "price": best_price
            })
        return result

    _prompt_price_detail = (
        "仅输出JSON对象（不含任何说明文字），识别充电价格详情图片中的分时段电价信息。图片中通常包含“充电时段”、“单价”、“电费”、“服务费”等列。\n"
        "请识别出所有的时段行，返回一个列表，每个元素包含：\n"
        "1. start: 开始时间 (HH:MM)；\n"
        "2. end: 结束时间 (HH:MM)；\n"
        "3. price: 总单价 (数字，如 0.6100)；\n"
        "4. ele_fee: 电费 (数字)；\n"
        "5. ser_fee: 服务费 (数字)。\n"
        "\n"
        "注意：\n"
        "- 只需识别“直流桩”或当前选中的桩型下的数据。\n"
        "- 严禁输出 Markdown 代码块标签，严格返回纯JSON对象。"
    )

    @classmethod
    async def get_price_detail_from_image(cls, image_path: str):
        """
        使用 VL 模型从三级价格详情页截图中识别分时段电价
        """
        if not os.path.exists(image_path):
            logger.error(f"Image not found: {image_path}")
            return None

        with open(image_path, "rb") as image_file:
            encoded_image = base64.b64encode(image_file.read()).decode("utf-8")

        try:
            response = await cls._client.chat.completions.create(
                model=VL_MODEL_NAME,
                messages=[
                    {
                        "role": "user",
                        "content": [
                            {"type": "text", "text": cls._prompt_price_detail},
                            {
                                "type": "image_url",
                                "image_url": {"url": f"data:image/jpeg;base64,{encoded_image}"},
                            },
                        ],
                    }
                ],
                max_tokens=2000,
                temperature=0.01
            )

            content = response.choices[0].message.content
            logger.info(f"VL Price Detail Response: {content}")

            json_str = cls._extract_json(content)
            data = json.loads(json_str)
            
            # 兼容性处理：如果返回的是对象且包含列表字段，提取列表
            if isinstance(data, dict):
                for key in ["price_list", "prices", "schedule", "data", "items"]:
                    if key in data and isinstance(data[key], list):
                        return data[key]
                
                # 如果字典本身就是一个价格项（包含 start 和 price），将其包装成列表
                if "start" in data and ("price" in data or "total_price" in data):
                    return [data]
            
            # 如果是列表，直接返回
            if isinstance(data, list):
                return data
            
            # 兜底：确保返回的是列表
            return [data] if data else []

        except Exception as e:
            logger.error(f"Error calling VL model for price detail: {e}")
            return None

    @classmethod
    async def get_station_detail_from_image(cls, image_path: str):
        """
        使用 VL 模型从详情页截图中识别地址、名称和电价
        """
        if not os.path.exists(image_path):
            logger.error(f"Image not found: {image_path}")
            return None

        with open(image_path, "rb") as image_file:
            encoded_image = base64.b64encode(image_file.read()).decode("utf-8")

        try:
            response = await cls._client.chat.completions.create(
                model=VL_MODEL_NAME,
                messages=[
                    {
                        "role": "user",
                        "content": [
                            {"type": "text", "text": cls._prompt_detail},
                            {
                                "type": "image_url",
                                "image_url": {"url": f"data:image/jpeg;base64,{encoded_image}"},
                            },
                        ],
                    }
                ],
                max_tokens=1000,
                temperature=0.01
            )

            content = response.choices[0].message.content
            logger.info(f"VL Detail Response: {content}")

            json_str = cls._extract_json(content)
            detail = json.loads(json_str)
            
            return detail

        except Exception as e:
            logger.error(f"Error calling VL model for detail: {e}")
            return None

    @staticmethod
    async def recognize_card_with_vl(image_path: str):
        """
        对卡片切片进行 VL 识别
        """
        if not os.path.exists(image_path):
            return None

        with open(image_path, "rb") as f:
            encoded = base64.b64encode(f.read()).decode("utf-8")
            img_input = {"url": f"data:image/jpeg;base64,{encoded}"}

        prompt = (
            "请识别图片中的充电站卡片信息，并以 JSON 格式输出：\n"
            "{\n"
            "  \"station_name\": \"场站名称\",\n"
            "  \"price\": \"价格(数字)\",\n"
            "  \"piles\": \"空闲/总数(如 4/4)\",\n"
            "  \"distance\": \"距离(如 2.5km)\",\n"
            "  \"parking\": \"停车描述\"\n"
            "}\n"
            "只输出 JSON，不要有其他文字。"
        )

        try:
            resp = await ReadImageKit._client.chat.completions.create(
                model=VL_MODEL_NAME,
                messages=[
                    {
                        "role": "user",
                        "content": [
                            {"type": "text", "text": prompt},
                            {"type": "image_url", "image_url": img_input}
                        ]
                    }
                ],
                max_tokens=300,
                temperature=0.01
            )
            content = resp.choices[0].message.content.strip()
            raw = ReadImageKit._extract_json(content)
            return json.loads(raw)
        except Exception as e:
            logger.error(f"VL recognize_card_with_vl error: {e}")
            return None

    @staticmethod
    async def get_stations_hybrid(image_path, device_info=None):
        """
        混合识别模式：图形学切片 + VL 模型识别
        """
        if device_info is None:
            device_info = ReadImageKit._FALLBACK_DEVICE_INFO

        # 1. 使用 Kit 中的图形学算法识别卡片区域
        json_data = Kit.crop_cards_from_image(image_path, save_debug=True)
        if not json_data or not json_data.get("cards"):
            logger.warning("No cards detected by graphical slicing.")
            return []

        # 2. 对每个卡片区域进行 VL 识别
        tasks = []
        card_metas = []
        
        # 为了使用 VL，我们需要将每个卡片切出来保存到临时文件
        temp_dir = os.path.join(os.path.dirname(image_path), "Patches")
        os.makedirs(temp_dir, exist_ok=True)
        
        img = Image.open(image_path).convert("RGB")
        
        for card in json_data["cards"]:
            rect = card["rect"] # [x1, y1, x2, y2]
            patch_path = os.path.join(temp_dir, f"patch_{card['id']}_{uuid.uuid4().hex[:8]}.jpg")
            
            # 裁剪并保存
            patch = img.crop((rect[0], rect[1], rect[2], rect[3]))
            patch.save(patch_path)
            
            card_metas.append({
                "meta": card,
                "patch_path": patch_path
            })
            tasks.append(ReadImageKit.recognize_card_with_vl(patch_path))

        # 并行执行 VL 识别
        results = await asyncio.gather(*tasks)
        
        final_stations = []
        for i, res in enumerate(results):
            meta = card_metas[i]["meta"]
            patch_path = card_metas[i]["patch_path"]
            
            if res and isinstance(res, dict) and res.get("station_name"):
                # 注入点击坐标和原始区域信息
                res["uia_center_x"] = meta["click_point"][0]
                res["uia_center_y"] = meta["click_point"][1]
                res["rect"] = meta["rect"]
                
                # 转换 bounds 到 0-1000 空间
                w, h = img.size
                rect = meta["rect"]
                res["bounds"] = [
                    int(rect[0] * 1000 / w),
                    int(rect[1] * 1000 / h),
                    int(rect[2] * 1000 / w),
                    int(rect[3] * 1000 / h)
                ]
                
                final_stations.append(res)
                logger.info(f"卡片 {meta['id']} 识别成功: {res['station_name']}")
            else:
                logger.warning(f"卡片 {meta['id']} 识别失败或无名称")
            
            # 清理临时切片文件
            try:
                if os.path.exists(patch_path):
                    os.remove(patch_path)
            except:
                pass

        return final_stations

    @classmethod
    async def get_stations_from_image(cls, image_path: str, device_info=None):
        """
        使用 Qwen-VL 模型从截图中识别充电站列表
        """
        if device_info is None:
            device_info = cls._FALLBACK_DEVICE_INFO

        if not os.path.exists(image_path):
            logger.error(f"Image not found: {image_path}")
            return []

        # 将图片转换为 Base64
        with open(image_path, "rb") as image_file:
            encoded_image = base64.b64encode(image_file.read()).decode("utf-8")

        try:
            response = await cls._client.chat.completions.create(
                model=VL_MODEL_NAME,
                messages=[
                    {
                        "role": "user",
                        "content": [
                            {"type": "text", "text": cls._prompt},
                            {
                                "type": "image_url",
                                "image_url": {"url": f"data:image/jpeg;base64,{encoded_image}"},
                            },
                        ],
                    }
                ],
                max_tokens=2000,
                temperature=0.01
            )

            content = response.choices[0].message.content
            logger.info(f"VL Model Response: {content}")

            json_str = cls._extract_json(content)
            stations = json.loads(json_str)

            # 后处理：如果 bounds 是归一化的，则转换为像素坐标（如果需要）
            # 或者如果 bounds 是 0-1000 的，则保持原样或按需转换
            return stations

        except Exception as e:
            logger.error(f"Error calling VL model: {e}")
            return []