#!/usr/bin/env python
# -*- coding: UTF-8 -*-
"""
@Project ：django_base_ai
@File    ：ai_chat.py
@Desc    ：AI 会话 ViewSet（django.mdc：有对应 Model 用 CustomModelViewSet），支持会话 CRUD 与发消息
"""

# ---------- 标准库 ----------
import base64  # 图片转 base64 用于多模态 content
import json  # 请求体解析、SSE 序列化、PPT 大纲 JSON 解析
import logging  # 本模块日志
import os  # 本地附件路径拼接与文件检查
import re  # html_url 中 slides/ 路径解析
import uuid  # 生成 PPT HTML 文件名
from typing import Any, Iterator  # 类型注解

# ---------- 第三方 ----------
import requests  # 按 file_urls 拉取远程图片

# ---------- Django ----------
from django.conf import settings  # MEDIA_URL/MEDIA_ROOT、AI_DEFAULT_PROVIDER
from django.db.models import Q  # 查询条件（本文件未直接使用，保留供扩展）
from django.core.files.uploadedfile import UploadedFile  # 判断上传文件类型
from django.http import StreamingHttpResponse  # 流式 SSE 响应
# ---------- DRF ----------
from rest_framework import serializers  # SerializerMethodField 等
from rest_framework.decorators import action  # ViewSet 自定义 action
from rest_framework.permissions import IsAuthenticated  # 需登录
from rest_framework.request import Request  # 请求类型
from rest_framework.response import Response  # 响应类型

# ---------- 本包：模型与 AI 工具 ----------
from django_base_ai.system.models import AIChatMessage, AIChatSession  # 会话与消息表
from django_base_ai.utils.ai import ChatMessage, ChatSession  # 对话与客户端（get_client 等由 ai_helpers 使用）
from django_base_ai.utils.ai.wanxiang import generate_image as wanxiang_generate_image  # 通义万相文生图
from django_base_ai.utils.json_response import DetailResponse, ErrorResponse  # 统一响应格式
from django_base_ai.system.views.ai_helpers import AIViewsHelper  # 共用：parse_body、get_client、sse_data、sse_event
from django_base_ai.system.views.ai_ppt import AIPptMixin  # AI PPT 大纲/HTML/布局/pptx
from django_base_ai.system.views.ai_docx import AIDocxMixin  # 段落处理 text_refine 等
from django_base_ai.utils.xss_cleaned import cleaned_string  # 用户输入 XSS 清洗
from django_base_ai.utils.serializers import CustomModelSerializer  # 基类序列化器
from django_base_ai.utils.viewset import CustomModelViewSet  # 基类 ViewSet

# 本模块日志，便于排查 AI 调用、附件、落库等问题
logger = logging.getLogger(__name__)


# 附件：仅支持图片，单文件最大 5MB，用于多模态（视觉）模型
ALLOWED_IMAGE_CONTENT_TYPES = {"image/jpeg", "image/png", "image/gif", "image/webp"}
MAX_ATTACHMENT_SIZE = 5 * 1024 * 1024  # 5MB


def _parse_body_and_files(request: Request) -> tuple[dict[str, Any], list[Any]]:
    """
    解析请求体与上传文件。返回 (body, files)。
    body 含 message、session_id、stream、file_urls 等；files 为 FILES 中的文件列表（支持 key 为 files 或 file）。
    """
    body = AIViewsHelper.parse_body(request)
    files: list[Any] = []
    flist = getattr(request, "FILES", None)
    if flist:
        # 兼容前端传 files[] 或 file[]
        files = list(flist.getlist("files") or flist.getlist("file") or [])
    return body, files


def _uploaded_files_to_image_parts(files: list[Any]) -> list[dict[str, Any]]:
    """
    将上传的图片转为 OpenAI 多模态 content 中的 image_url 部分。
    超出大小或非图片类型跳过并打日志；格式为 data:image/xxx;base64,...
    """
    parts: list[dict[str, Any]] = []
    for f in files:
        if not isinstance(f, UploadedFile):
            continue
        if f.size > MAX_ATTACHMENT_SIZE:
            logger.warning("Attachment %s exceeds max size %s", getattr(f, "name", ""), MAX_ATTACHMENT_SIZE)
            continue
        # 取 MIME 主类型，忽略 charset 等参数
        content_type = (getattr(f, "content_type", "") or "").split(";")[0].strip().lower()
        if content_type not in ALLOWED_IMAGE_CONTENT_TYPES:
            logger.warning("Attachment %s type %s not allowed", getattr(f, "name", ""), content_type)
            continue
        try:
            raw = f.read()
            b64 = base64.b64encode(raw).decode("ascii")
            data_url = f"data:{content_type};base64,{b64}"
            # OpenAI 多模态格式：type=image_url, image_url.url=data URL
            parts.append({"type": "image_url", "image_url": {"url": data_url}})
        except Exception as e:
            logger.warning("Failed to read attachment %s: %s", getattr(f, "name", ""), e)
    return parts


def _file_urls_to_image_parts(request: Request, file_urls: Any) -> list[dict[str, Any]]:
    """
    将「统一上传接口 /base/api/system/file/ 返回的文件 URL」拉取并转为多模态 image_url 部分。
    file_urls 为字符串或列表，每个元素为上传接口返回的 url（可相对路径如 /media/xxx 或绝对 URL）。
    优先从本地 MEDIA_ROOT 读文件，否则用 HTTP GET 拉取。
    """
    parts: list[dict[str, Any]] = []
    if not file_urls:
        return parts
    if isinstance(file_urls, str):
        file_urls = [file_urls]
    if not isinstance(file_urls, list):
        return parts
    media_url = (getattr(settings, "MEDIA_URL", "") or "media/").strip().rstrip("/")
    media_root = getattr(settings, "MEDIA_ROOT", "") or "media"
    for raw_url in file_urls:
        if not raw_url or not isinstance(raw_url, str):
            continue
        url = raw_url.strip()
        if not url:
            continue
        content_type: str | None = None
        raw: bytes | None = None
        # 相对路径或含 MEDIA 的路径：尝试从本地 MEDIA_ROOT 读取，避免走外网
        if url.startswith("/") or url.startswith(media_url):
            path = url.lstrip("/")
            if path.startswith(media_url):
                path = path[len(media_url) :].lstrip("/")
            local_path = os.path.join(media_root, path) if not os.path.isabs(path) else path
            if os.path.isfile(local_path) and os.path.getsize(local_path) <= MAX_ATTACHMENT_SIZE:
                try:
                    with open(local_path, "rb") as f:
                        raw = f.read()
                    ext = os.path.splitext(local_path)[1].lower()
                    content_type = {
                        ".jpg": "image/jpeg", ".jpeg": "image/jpeg", ".png": "image/png",
                        ".gif": "image/gif", ".webp": "image/webp",
                    }.get(ext)
                except Exception as e:
                    logger.warning("Read local file %s failed: %s", local_path, e)
        # 非本地或本地读取失败：按绝对 URL 或相对 URL 用 HTTP 拉取
        if raw is None:
            fetch_url = request.build_absolute_uri(url) if url.startswith("/") else url
            try:
                r = requests.get(fetch_url, timeout=15)
                r.raise_for_status()
                if len(r.content) > MAX_ATTACHMENT_SIZE:
                    logger.warning("File %s exceeds max size", fetch_url)
                    continue
                raw = r.content
                content_type = (r.headers.get("Content-Type") or "").split(";")[0].strip().lower()
            except Exception as e:
                logger.warning("Fetch file_url %s failed: %s", fetch_url, e)
                continue
        if not raw or content_type not in ALLOWED_IMAGE_CONTENT_TYPES:
            continue
        try:
            b64 = base64.b64encode(raw).decode("ascii")
            data_url = f"data:{content_type};base64,{b64}"
            parts.append({"type": "image_url", "image_url": {"url": data_url}})
        except Exception as e:
            logger.warning("Encode file_url failed: %s", e)
    return parts


def _user_message_text_for_db(content: Any) -> str:
    """
    从本轮用户消息 content（str 或多模态 list）中取出纯文本，用于落库。
    多模态时取 type=text 的 text；若仅有图片则返回 "[含n张图片]" 占位。
    """
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        for part in content:
            if isinstance(part, dict) and part.get("type") == "text" and part.get("text"):
                return part["text"]
        n = sum(1 for p in content if isinstance(p, dict) and p.get("type") == "image_url")
        return f"[含{n}张图片]" if n else ""
    return ""


def _load_or_create_session(
    request: Request, body: dict[str, Any], client: Any, first_message: str
) -> tuple[AIChatSession | None, str | None]:
    """
    按 session_id 加载已有会话（并校验归属），否则创建新会话。返回 (session, error_msg)。
    有 session_id 时校验：会话存在且归属当前用户；新建时 title 取首条消息前 50 字。
    """
    session_id = body.get("session_id")
    model = body.get("model")
    default_model = getattr(client, "default_model", "") or ""

    if session_id:
        session_obj = AIChatSession.objects.filter(id=session_id).first()
        if not session_obj:
            return None, "会话不存在"
        # 有归属的会话仅创建者可操作
        if session_obj.user_id and request.user and session_obj.user_id != request.user.id:
            return None, "无权限操作该会话"
        return session_obj, None

    # 新建会话：标题用首条消息前 50 字，后续流式结束时可能再更新
    session_obj = AIChatSession(
        user=request.user,
        title=(first_message[:50] if first_message else ""),
        model_name=model or default_model,
    )
    session_obj.creator = str(request.user.id)
    session_obj.modifier = str(request.user.id)
    session_obj.save()
    return session_obj, None


def _content_for_db(content: str) -> str:
    """
    将内容转为 MySQL utf8mb3 可存格式：去掉 4 字节字符（如 emoji），避免 Incorrect string value。
    若表已是 utf8mb4，可不再调用此函数。
    """
    if not content:
        return content
    return "".join(c if ord(c) <= 0xFFFF else "" for c in content)


def _user_content_for_db(text: str) -> str:
    """
    用户上送内容落库前做 XSS 清洗再转 utf8，用于 AIChatMessage 用户消息及会话 title。
    先 cleaned_string(escape_html=True)，再 _content_for_db 去 4 字节字符。
    """
    if not text:
        return text
    return _content_for_db(cleaned_string(text, escape_html=True))


def _session_messages_for_api(session: AIChatSession) -> list[ChatMessage]:
    """
    将会话内历史消息转为 ChatMessage 列表，按 id 升序。
    供调用 client.chat / client.chat_stream 时作为历史上下文。
    """
    return [
        ChatMessage(role=m.role, content=m.content or "")
        for m in AIChatMessage.objects.filter(session=session).order_by("id")
    ]


def _sse_chat_stream(
    request: Request,
    session_obj: AIChatSession,
    message_content: Any,
    messages_for_api: list[ChatMessage],
    client: Any,
    model: str | None,
) -> Iterator[str]:
    """
    流式对话 SSE 生成器。message_content 为 str 或多模态 list（含图片时）。
    每轮 yield 一条 data: {json}\\n\\n；结束时该条含 done、usage、session_id、message_id，并落库用户消息与助手消息。
    """
    stream_fn = getattr(client, "chat_stream", None)
    if not stream_fn:
        yield AIViewsHelper.sse_data({"error": "当前供应商不支持流式输出"})
        return
    messages = messages_for_api + [ChatMessage(role="user", content=message_content)]
    full_content: list[str] = []  # 累积助手回复文本，done 时一次性落库
    message_id = None
    text_for_db = _user_message_text_for_db(message_content)
    try:
        for chunk in stream_fn(messages, model=model):
            if chunk.get("content"):
                full_content.append(chunk["content"])
            # 流结束：落库用户消息 + 助手消息，并补全 session_id / message_id 供前端
            if chunk.get("done"):
                uid = str(getattr(request.user, "id", None) or "")
                try:
                    AIChatMessage.objects.create(
                        session=session_obj,
                        role="user",
                        content=_user_content_for_db(text_for_db),
                        creator=uid,
                    )
                    full_str = "".join(full_content)
                    msg_obj = AIChatMessage.objects.create(
                        session=session_obj,
                        role="assistant",
                        content=_content_for_db(full_str),
                        model_name=chunk.get("model") or "",
                        usage=chunk.get("usage") or {},
                        creator=uid,
                    )
                    message_id = msg_obj.id
                    if not session_obj.title and text_for_db:
                        session_obj.title = _user_content_for_db(text_for_db)[:50]
                        session_obj.save(update_fields=["title", "update_datetime"])
                except Exception as e:
                    logger.exception("Save AI stream message failed: %s", e)
                chunk["session_id"] = session_obj.id
                chunk["message_id"] = message_id
            yield AIViewsHelper.sse_data(chunk)
    except Exception as e:
        logger.exception("AI chat stream failed")
        yield AIViewsHelper.sse_data({"error": str(e), "done": True})


def _sse_chat_stream_bytes(
    request: Request,
    session_obj: AIChatSession,
    message_content: Any,
    messages_for_api: list[ChatMessage],
    client: Any,
    model: str | None,
):
    """
    与 _sse_chat_stream 相同，但逐段 yield UTF-8 bytes。
    StreamingHttpResponse 使用迭代器时需 yield bytes，否则中文可能乱码。
    """
    for chunk in _sse_chat_stream(
        request, session_obj, message_content, messages_for_api, client, model
    ):
        yield chunk.encode("utf-8")


# ---------- Serializer（django.mdc：有对应 Model 时配套 List/Detail 与 Create/Update）----------
class AIChatSessionSerializer(CustomModelSerializer):
    """列表/详情用，只读 id，其余字段可序列化/反序列化。"""

    class Meta:
        model = AIChatSession
        fields = "__all__"
        read_only_fields = ["id"]


class AIChatSessionThumbnailSerializer(CustomModelSerializer):
    """
    用于 list 接口：返回用户发起的所有会话的缩略信息，仅含列表展示所需字段 + attachment_type。
    前端可按 create_datetime / update_datetime 分组（如昨天、7天内、30天内）。
    """
    attachment_type = serializers.SerializerMethodField(read_only=True)

    def get_attachment_type(self, obj) -> str:
        """附件类型：当前消息仅存文本，固定返回 text；若后续落库附件标记可改为 image 等。"""
        return "text"

    class Meta:
        model = AIChatSession
        fields = ["id", "title", "create_datetime", "update_datetime", "model_name", "attachment_type"]
        read_only_fields = ["id", "title", "create_datetime", "update_datetime", "model_name", "attachment_type"]


class AIChatSessionCreateUpdateSerializer(CustomModelSerializer):
    """创建/更新会话时使用，字段全开放。"""

    class Meta:
        model = AIChatSession
        fields = "__all__"


class AIChatSessionViewSet(AIPptMixin, AIDocxMixin, CustomModelViewSet):
    """
    AI 会话：list/retrieve/create/update/destroy 会话；chat action 发消息并落库 AIChatMessage。
    PPT 来自 AIPptMixin，段落处理 text_refine 来自 AIDocxMixin；依赖本类提供 _parse_body、_get_client、_sse_event。
    list 仅返回当前用户发起的会话，使用缩略序列化（id、title、时间、model_name、attachment_type）。
    会话隔离：session_id 为表主键 id，通过「仅返回当前用户会话」+「发消息时校验归属」保证不会串会话。
    """
    queryset = AIChatSession.objects.all()  # 实际 list 时由基类/过滤仅当前用户
    serializer_class = AIChatSessionSerializer  # retrieve 等用
    list_serializer_class = AIChatSessionThumbnailSerializer  # list 用缩略
    create_serializer_class = AIChatSessionCreateUpdateSerializer
    update_serializer_class = AIChatSessionCreateUpdateSerializer
    search_fields = ["title", "model_name"]  # 列表搜索

    # 供 AIPptMixin / AIDocxMixin / AIExcelMixin 使用（共用逻辑在 ai_helpers.AIViewsHelper）
    def _parse_body(self, request: Request) -> dict[str, Any]:
        return AIViewsHelper.parse_body(request)

    def _get_client(self, body: dict[str, Any]) -> tuple[Any, str | None]:
        return AIViewsHelper.get_client(body)

    def _sse_event(self, event_name: str, data: dict[str, Any]) -> str:
        return AIViewsHelper.sse_event(event_name, data)

    @action(detail=False, methods=["post"], permission_classes=[IsAuthenticated])
    def chat(self, request: Request) -> Response:
        """
        发消息并返回 AI 回复，支持新会话、续写历史、上传附件（图片）。
        请求体: message（必填）, stream, session_id, provider, model 等；
               图片可二选一或同时使用：① multipart 的 files/file；② file_urls 或 attachment_urls（数组，为统一上传接口 /base/api/system/file/ 返回的 url）。
        图片仅支持 jpeg/png/gif/webp，单文件最大 5MB，需模型支持视觉（多模态）。
        stream=0 时响应: { content, usage, model, session_id, message_id }；
        stream=1 时响应: text/event-stream。
        """
        # 1) 解析请求体与上传文件
        body, uploaded_files = _parse_body_and_files(request)
        message = str((body.get("message") or "")).strip()
        if not message:
            return ErrorResponse(msg="message 不能为空")

        # 2) 图片来源：本次上传的 files + 统一上传接口返回的 file_urls（/base/api/system/file/ 返回的 url）
        image_parts = _uploaded_files_to_image_parts(uploaded_files)
        file_urls = body.get("file_urls") or body.get("attachment_urls")
        image_parts = image_parts + _file_urls_to_image_parts(request, file_urls)
        # 构建本轮用户消息 content：纯文本或 文本+图片（多模态 list）
        if image_parts:
            user_message_content: Any = [{"type": "text", "text": message}] + image_parts
        else:
            user_message_content = message

        # 3) 流式/非流式：stream=1 为 SSE 流式
        stream_mode = body.get("stream", 0)
        try:
            stream_mode = int(stream_mode) if stream_mode is not None else 0
        except (TypeError, ValueError):
            stream_mode = 0
        use_stream = stream_mode == 1

        # 4) 创建 AI 客户端（use_settings 或 body 内 api_key）
        client, err = AIViewsHelper.get_client(body)
        if err:
            return ErrorResponse(msg=err)

        # 5) 按 session_id 加载或新建会话，并校验归属
        session_obj, err = _load_or_create_session(
            request, body, client, message
        )
        if err:
            status = 403 if "无权限" in err else 400
            return ErrorResponse(msg=err, status=status)

        messages_for_api = _session_messages_for_api(session_obj)
        model = body.get("model")

        # 6) 流式：返回 SSE，在生成器内落库
        if use_stream:
            resp = StreamingHttpResponse(
                _sse_chat_stream_bytes(
                    request, session_obj, user_message_content, messages_for_api, client, model
                ),
                content_type="text/event-stream; charset=utf-8",
            )
            resp["Cache-Control"] = "no-cache"
            resp["X-Accel-Buffering"] = "no"
            resp["Connection"] = "keep-alive"
            return resp

        # 7) 非流式：有附件（多模态 list）时直接 client.chat；无附件用 ChatSession.send 走历史轮数控制
        if isinstance(user_message_content, list):
            messages = messages_for_api + [ChatMessage(role="user", content=user_message_content)]
            try:
                result = client.chat(messages, model=model)
            except Exception as e:
                logger.exception("AI chat request failed")
                return ErrorResponse(msg=f"AI 调用失败: {e}")
        else:
            try:
                chat_session = ChatSession(client, max_history_turns=20)
                chat_session._messages = messages_for_api
                result = chat_session.send(user_message_content, model=model)
            except Exception as e:
                logger.exception("AI chat request failed")
                return ErrorResponse(msg=f"AI 调用失败: {e}")

        # 8) 落库用户消息与助手消息，并补全会话 title（若为空）
        text_for_db = _user_message_text_for_db(user_message_content)
        user_id = getattr(request.user, "id", None)
        uid = str(user_id) if user_id else None
        message_id = None
        try:
            AIChatMessage.objects.create(
                session=session_obj,
                role="user",
                content=_user_content_for_db(text_for_db),
                creator=uid,
            )
            msg_obj = AIChatMessage.objects.create(
                session=session_obj,
                role="assistant",
                content=_content_for_db(result.content),
                model_name=result.model or "",
                usage=result.usage or {},
                creator=uid,
            )
            message_id = msg_obj.id
            if not session_obj.title and text_for_db:
                session_obj.title = _user_content_for_db(text_for_db)[:50]
                session_obj.save(update_fields=["title", "update_datetime"])
        except Exception as e:
            logger.exception("Save AI message failed")
            return ErrorResponse(msg="保存会话消息失败，请重试", status=500)

        return DetailResponse(
            data={
                "content": result.content,
                "usage": result.usage or {},
                "model": result.model or "",
                "session_id": session_obj.id,
                "message_id": message_id,
            }
        )
