from appPublic.timeUtils import curDateString
from filemgr.filemgr import FileMgr
from rag.uapi_service import APIService
from appPublic.registerfunction import RegisterFunction
from appPublic.log import debug, error, info
from sqlor.dbpools import DBPools
import asyncio
import aiohttp
from langchain_core.documents import Document
from langchain_text_splitters import RecursiveCharacterTextSplitter
import os
import re
import time
import uuid
from datetime import datetime
import traceback
from filetxt.loader import fileloader,File2Text
from ahserver.serverenv import get_serverenv
from typing import List, Dict, Any
from rag.service_opts import get_service_params, sor_get_service_params, sor_get_embedding_mode, get_embedding_mode
from rag.fileprocess import extract_images_from_file
from rag.rag_operations import RagOperations
import json
from rag.transaction_manager import TransactionContext
from dataclasses import dataclass
from enum import Enum
import base64
from pathlib import Path

class RagFileMgr(FileMgr):
	def __init__(self, fiid):
		super().__init__(fiid)
		self.rag_ops = RagOperations()

	async def get_folder_ownerid(self, sor):
		fiid = self.fiid
		recs = await sor.R('kdb', {'id': self.fiid})
		if len(recs) > 0:
			return recs[0].orgid
		return None

	async def get_organization_quota(self, sor, orgid):
		sql = """select a.* from ragquota a, kdb b
where a.orgid = b.orgid
	and b.id = ${id}$
	and ${today}$ >= a.enabled_date
	and ${today}$ < a.expired_date
"""
		recs = await sor.sqlExe(sql, {
			 'id': self.fiid,
			 'today': curDateString()
		})
		if len(recs) > 0:
			r = recs[0]
			return r.quota, r.expired_date
		return None, None

	async def file_to_base64(self,path: str) -> str:
		with open(path, "rb") as f:
			return base64.b64encode(f.read()).decode("utf-8")

	async def file_uploaded(self, request, ns, userid):
		"""将文档插入 Milvus 并抽取三元组到 Neo4j"""
		debug(f'Received ns: {ns=}')
		env = request._run_ns
		realpath = ns.get('realpath', '')
		fiid = ns.get('fiid', '')
		id = ns.get('id', '')
		orgid = ns.get('ownerid', '')
		db_type = ''

		debug(
			f'Inserting document: file_path={realpath}, userid={orgid}, db_type={db_type}, knowledge_base_id={fiid}, document_id={id}')

		timings = {}
		start_total = time.time()
		result = {
			"status": "error",
			"userid": orgid,
			"document_id": id,
			"collection_name": "ragdb",
			"timings": timings,
			"message": "",
			"status_code": 400
		}

		# 初始化回滚上下文
		rollback_context = {
			"request": request,
			"userid": userid,
			"service_params": None  # 在 try 块中设置
		}

		async with TransactionContext(f"file_upload_{id}") as transaction_mgr:
			# 将 rollback_context 绑定到 TransactionContext
			transaction_mgr.transaction_context = rollback_context
			try:
				# 验证必填字段
				if not orgid or not fiid or not id:
					raise ValueError("orgid、fiid 和 id 不能为空")
				if len(orgid) > 32 or len(fiid) > 255:
					raise ValueError("orgid 或 fiid 的长度超出限制")
				if not os.path.exists(realpath):
					raise ValueError(f"文件 {realpath} 不存在")

				# 获取服务参数
				service_params = await get_service_params(orgid)
				debug(f"服务参数是：{service_params}")
				if not service_params:
					raise ValueError("无法获取服务参数")
				rollback_context["service_params"] = service_params

				#获取嵌入模式
				embedding_mode = await get_embedding_mode(orgid)
				debug(f"检测到 embedding_mode = {embedding_mode}（0=文本, 1=多模态）")

				# 加载和分片文档
				chunks = await self.rag_ops.load_and_chunk_document(
					realpath, timings, transaction_mgr=transaction_mgr
				)

				text_embeddings = None
				multi_results = None
				image_paths = []

				if embedding_mode == 1:
					inputs = []
					# 文本
					for chunk in chunks:
						inputs.append({"type": "text", "content": chunk.page_content})

					debug("开始多模态图像抽取与嵌入")
					image_paths = extract_images_from_file(realpath)
					debug(f"从文档中抽取 {len(image_paths)} 张图像")

					if image_paths:
						for img_path in image_paths:
							try:
								# 1. 自动识别真实格式
								ext = Path(img_path).suffix.lower()
								if ext not in {".png", ".jpg", ".jpeg", ".webp", ".bmp"}:
									ext = ".jpg"

								mime_map = {
									".png": "image/png",
									".jpg": "image/jpeg",
									".jpeg": "image/jpeg",
									".webp": "image/webp",
									".bmp": "image/bmp"
								}
								mime_type = mime_map.get(ext, "image/jpeg")

								# # 2. 智能压缩（>1MB 才压缩，节省 70% 流量）
								# img = Image.open(img_path).convert("RGB")
								# if os.path.getsize(img_path) > 1024 * 1024:  # >1MB
								# 	buffer = BytesIO()
								# 	img.save(buffer, format="JPEG", quality=85, optimize=True)
								# 	b64 = base64.b64encode(buffer.getvalue()).decode()
								# 	data_uri = f"data:image/jpeg;base64,{b64}"
								# else:
								b64 = await self.file_to_base64(img_path)
								data_uri = f"data:{mime_type};base64,{b64}"

								inputs.append({
									"type": "image",
									"data": data_uri
								})
								debug(f"已添加图像（{mime_type}, {len(b64) / 1024:.1f}KB）: {Path(img_path).name}")

							except Exception as e:
								debug(f"图像处理失败，跳过: {img_path} → {e}")
								# 即使失败也加个占位，防止顺序错乱
								inputs.append({
									"type": "image",
									"data": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNkYAAAAAYAAjCB0C8AAAAASUVORK5CYII="
								})

					debug(f"混排输入总数: {len(inputs)}（文本 {len(chunks)} + 图像 {len(image_paths)}）")

					multi_results = await self.rag_ops.generate_multi_embeddings(
						request=request,
						inputs=inputs,
						service_params=service_params,
						userid=userid,
						timings=timings,
						transaction_mgr=transaction_mgr
					)
					debug(f"多模态嵌入成功，返回 {len(multi_results)} 条结果")
				else:
					# 生成嵌入向量
					debug("【纯文本模式】使用 BGE 嵌入")
					text_embeddings = await self.rag_ops.generate_embeddings(
						request, chunks, service_params, userid, timings, transaction_mgr=transaction_mgr
					)
					debug(f"BGE 嵌入完成: {len(text_embeddings)} 条")

				inserted = await self.rag_ops.insert_all_vectors(
					request=request,
					text_chunks=chunks,
					realpath=realpath,
					orgid=orgid,
					fiid=fiid,
					document_id=id,
					service_params=service_params,
					userid=userid,
					db_type=db_type,
					timings=timings,
					img_paths=image_paths,
					text_embeddings=text_embeddings,
					multi_results=multi_results,
					transaction_mgr=transaction_mgr
				)
				debug(f"统一插入: 文本 {inserted['text']}, 图像 {inserted['image']}, 人脸 {inserted['face']}")

				# 抽取三元组
				triples = await self.rag_ops.extract_triples(
					request, chunks, service_params, userid, timings, transaction_mgr=transaction_mgr
				)

				# 插入 Neo4j
				await self.rag_ops.insert_to_graph_db(
					request, triples, id, fiid, orgid, service_params, userid, timings, transaction_mgr=transaction_mgr
				)

				timings["total"] = time.time() - start_total
				result.update({
					"status": "success",
					"unique_triples": triples,
					"message": f"文件 {realpath} 成功嵌入并处理三元组",
					"status_code": 200
				})
				debug(f"总耗时: {timings['total']:.2f} 秒")

			except Exception as e:
				error(f"插入文档失败: {str(e)}, 堆栈: {traceback.format_exc()}")
				timings["total"] = time.time() - start_total
				result.update({
					"message": f"插入文档失败: {str(e)}",
					"timings": timings
				})
				raise ValueError(str(e)) from e
		debug(f"最终结果是：{result}")
		return result

	async def file_deleted(self, request, recs, userid):
		"""删除用户指定文件数据，包括 Milvus 和 Neo4j 中的记录"""
		if not isinstance(recs, list):
			recs = [recs]
		results = []
		total_nodes_deleted = 0
		total_rels_deleted = 0

		for rec in recs:
			id = rec.get('id', '')
			realpath = rec.get('realpath', '')
			fiid = rec.get('fiid', '')
			orgid = rec.get('ownerid', '')
			db_type = ''
			collection_name = "ragdb" if not db_type else f"ragdb_{db_type}"

			try:
				required_fields = ['id', 'realpath', 'fiid', 'ownerid']
				missing_fields = [field for field in required_fields if not rec.get(field, '')]
				if missing_fields:
					raise ValueError(f"缺少必填字段: {', '.join(missing_fields)}")

				service_params = await get_service_params(orgid)
				if not service_params:
					raise ValueError("无法获取服务参数")

				# 调用 Milvus 删除
				await self.rag_ops.delete_from_vector_db(request, orgid, realpath, fiid, id, service_params, userid, db_type)

				# 调用 Neo4j 删除
				neo4j_deleted_nodes = 0
				neo4j_deleted_rels = 0
				try:
					nodes_deleted, rels_deleted = await self.rag_ops.delete_from_graph_db(request, id, service_params, userid)
					neo4j_deleted_nodes += nodes_deleted
					neo4j_deleted_rels += rels_deleted
					total_nodes_deleted += nodes_deleted
					total_rels_deleted += rels_deleted
				except Exception as e:
					error(f"删除 document_id={id} 的 Neo4j 数据失败: {str(e)}")

				results.append({
					"status": "success",
					"collection_name": collection_name,
					"document_id": id,
					"message": f"成功删除文件 {realpath} 的 Milvus 记录，{neo4j_deleted_nodes} 个 Neo4j 节点，{neo4j_deleted_rels} 个 Neo4j 关系",
					"status_code": 200
				})

			except Exception as e:
				error(f"删除文档 {realpath} 失败: {str(e)}, 堆栈: {traceback.format_exc()}")
				results.append({
					"status": "error",
					"collection_name": collection_name,
					"document_id": id,
					"message": f"删除文档 {realpath} 失败: {str(e)}",
					"status_code": 400
				})

		return {
			"status": "success" if all(r["status"] == "success" for r in results) else "partial",
			"results": results,
			"total_nodes_deleted": total_nodes_deleted,
			"total_rels_deleted": total_rels_deleted,
			"message": f"处理 {len(recs)} 个文件，成功删除 {sum(1 for r in results if r['status'] == 'success')} 个",
			"status_code": 200 if all(r["status"] == "success" for r in results) else 207
		}