import asyncio, base64, datetime, inspect, json, logging, re, time
from concurrent.futures import ThreadPoolExecutor
from copy import deepcopy
from functools import partial
from typing import Any, Union, Tuple

from agent.component import component_class          # 动态组件工厂
from agent.component.base import ComponentBase         # 组件基类
from agent.dsl_migration import normalize_chunker_dsl  # DSL 版本迁移
from api.db.services.file_service import FileService   # 文件解析服务
from api.db.services.llm_service import LLMBundle      # LLM 调用封装
from api.db.services.task_service import has_canceled  # 取消状态查询
from api.db.joint_services.tenant_model_service import get_tenant_default_model_by_type
from common.constants import LLMType
from common.misc_utils import get_uuid, hash_str2int
from common.exceptions import TaskCanceledException
from rag.prompts.generator import chunks_format
from rag.utils.redis_conn import REDIS_CONN
from rag.utils.tts_cache import synthesize_with_cache

dsl = {
    "components": {        # 组件映射表，key 为组件实例ID
        "begin": {
            "obj": {       # 组件对象及其参数
                "component_name": "Begin",
                "params": {},
            },
            "downstream": ["answer_0"],   # 下游组件 ID 列表
            "upstream": [],               # 上游组件 ID 列表
        },
        "retrieval_0": {
            "obj": { "component_name": "Retrieval", "params": {} },
            "downstream": ["generate_0"],
            "upstream": ["answer_0"],
        },
        # ... 更多组件
    },
    "history": [],         # 对话历史
    "path": ["begin"],     # 当前执行路径（组件 ID 的有序列表）
    "retrieval": [         # 检索结果栈
        {"chunks": [], "doc_aggs": []}
    ],
    "globals": {           # 全局变量
        "sys.query": "",
        "sys.user_id": tenant_id,
        "sys.conversation_turns": 0,
        "sys.files": []
    }
}

def __init__(self, dsl: str, tenant_id=None, task_id=None, custom_header=None):
    self.path = []
    self.components = {}
    self.error = ""
    self.dsl = normalize_chunker_dsl(json.loads(dsl))  # ① DSL迁移
    self._tenant_id = tenant_id
    self.task_id = task_id if task_id else get_uuid()
    self.custom_header = custom_header
    self._thread_pool = ThreadPoolExecutor(max_workers=5)  # ② 线程池
    self.load()

def load(self):
    self.components = self.dsl["components"]
    cpn_nms = set([])
    for k, cpn in self.components.items():
        cpn_nms.add(cpn["obj"]["component_name"])
        param = component_class(cpn["obj"]["component_name"] + "Param")()
        param.update(cpn["obj"]["params"])
        try:
            param.check()
        except Exception as e:
            raise ValueError(self.get_component_name(k) + f": {e}")
        cpn["obj"] = component_class(cpn["obj"]["component_name"])(self, k, param)
    self.path = self.dsl["path"]

pat = re.compile(r"\{* *\{([a-zA-Z:0-9]+@[A-Za-z0-9_.-]+|sys\.[A-Za-z0-9_.]+|env\.[A-Za-z0-9_.]+)\} *\}*")

def get_variable_value(self, exp: str) -> Any:
    exp = exp.strip("{").strip("}").strip(" ").strip("{").strip("}")
    if exp.find("@") < 0:
        return self.globals[exp]                  # sys.* 或 env.* 全局变量
    cpn_id, var_nm = exp.split("@")
    cpn = self.get_component(cpn_id)
    parts = var_nm.split(".", 1)
    root_key = parts[0]
    rest = parts[1] if len(parts) > 1 else ""
    root_val = cpn["obj"].output(root_key)        # 调用组件的 output() 方法
    if not rest:
        return root_val
    return self.get_variable_param_value(root_val, rest)  # 深度取值

for key in path.split('.'):
    if isinstance(cur, dict):
        cur = cur.get(key)        # 字典按键访问
    elif isinstance(cur, (list, tuple)):
        cur = cur[int(key)]       # 列表按索引访问
    else:
        cur = getattr(cur, key, None)  # 对象按属性访问

def __str__(self):
    self.dsl["path"] = self.path
    self.dsl["task_id"] = self.task_id
    dsl = {"components": {}}
    for k in self.dsl.keys():
        if k in ["components"]:
            continue
        dsl[k] = deepcopy(self.dsl[k])
    for k, cpn in self.components.items():
        dsl["components"][k] = {}
        for c in cpn.keys():
            if c == "obj":
                dsl["components"][k][c] = json.loads(str(cpn["obj"]))
            else:
                dsl["components"][k][c] = deepcopy(cpn[c])
    return json.dumps(dsl, ensure_ascii=False)

def is_canceled(self) -> bool:
    return has_canceled(self.task_id)  # 检查 Redis key: {task_id}-cancel

def cancel_task(self) -> bool:
    REDIS_CONN.set(f"{self.task_id}-cancel", "x")  # 设置取消信号

def __init__(self, dsl: str, tenant_id=None, task_id=None, canvas_id=None, custom_header=None):
    self.globals = {
        "sys.query": "",
        "sys.user_id": tenant_id,
        "sys.conversation_turns": 0,
        "sys.files": [],
        "sys.history": [],
        "sys.date": datetime.datetime.now(datetime.timezone.utc).strftime("%Y-%m-%d %H:%M:%S")
    }
    self.variables = {}
    super().__init__(dsl, tenant_id, task_id, custom_header=custom_header)
    self._id = canvas_id

def load(self):
    super().load()
    self.history = self.dsl["history"]
    if "globals" in self.dsl:
        self.globals = self.dsl["globals"]
        # 确保必要的 key 存在
        if "sys.history" not in self.globals:
            self.globals["sys.history"] = []
        if "sys.date" not in self.globals:
            self.globals["sys.date"] = datetime.datetime.now(...)
    # ...
    self.retrieval = self.dsl["retrieval"]
    self.memory = self.dsl.get("memory", [])

def reset(self, mem=False):
    super().reset()
    if not mem:
        self.history = []
        self.retrieval = []
        self.memory = []
    for k in self.globals.keys():
        if k.startswith("sys."):
            # 按类型重置为默认值（"" / 0 / 0.0 / [] / {}）
        if k.startswith("env."):
            # 按用户定义的变量类型重置

def decorate(event, dt):
    nonlocal created_at
    return {
        "event": event,
        "message_id": self.message_id,
        "created_at": created_at,
        "task_id": self.task_id,
        "data": dt
    }

self.add_user_input(kwargs.get("query"))
# ...
for k in kwargs.keys():
    if k in ["query", "user_id", "files"] and kwargs[k]:
        if k == "files":
            self.globals[f"sys.{k}"] = await self.get_files_async(kwargs[k], ...)
        else:
            self.globals[f"sys.{k}"] = kwargs[k]
self.globals["sys.conversation_turns"] += 1

idx = len(self.path) - 1     # 当前批次的起始索引
while idx < len(self.path):  # 直到没有新节点加入
    to = len(self.path)       # 当前批次的结束索引
    # ① yield node_started 事件
    for i in range(idx, to):
        yield decorate("node_started", {...})

    # ② 批量并发执行
    await _run_batch(idx, to)

    # ③ post-processing：逐节点处理输出、错误、路径推进
    for i in range(idx, to):
        # ... 处理 Message 流式输出
        # ... 处理错误和异常跳转
        # ... 推进路径（下游节点加入 self.path）

    # ④ UserFillUp 检查
    if any(component_name == "userfillup" for ...):
        yield decorate("user_inputs", {...})
        return  # 等待用户输入
    idx = to

async def _run_batch(f, t):
    loop = asyncio.get_running_loop()
    tasks = []
    max_concurrency = 5
    sem = asyncio.Semaphore(max_concurrency)  # 限制最大并发数为 5

    async def _invoke_one(cpn_obj, sync_fn, call_kwargs, use_async):
        async with sem:
            if use_async:
                await cpn_obj.invoke_async(...)   # 异步调用
            else:
                await loop.run_in_executor(       # 线程池调用
                    self._thread_pool,
                    partial(sync_fn, ...)
                )
    # ...
    await asyncio.gather(*tasks)

for _, ele in cpn.get_input_elements().items():
    if isinstance(ele, dict) and ele.get("_cpn_id") and \
       ele.get("_cpn_id") not in self.path[:i]:
        self.path.pop(i)   # 依赖未就绪，从路径中移除
        t -= 1
        break

if cpn_obj.component_name.lower() == "message":
    if isinstance(cpn_obj.output("content"), partial):
        stream = cpn_obj.output("content")()
        # 处理流式输出，支持:
        # - <think></think> 标签 → start_to_think / end_to_think
        # - TTS 音频合成（每 16 个字符触发一次）
        # - audio_binary 字段传递 base64 编码的音频

# 迭代/循环内部项结束 → 回到父迭代器
if cpn_obj.component_name.lower() in ("iterationitem","loopitem") and cpn_obj.end():
    iter = cpn_obj.get_parent()
    _extend_path(self.get_component(cpn["parent_id"])["downstream"])

# 分类/分支组件 → 使用 _next 输出决定路径
elif cpn_obj.component_name.lower() in ["categorize", "switch"]:
    _extend_path(cpn_obj.output("_next"))

# 迭代/循环开始 → 进入内部第一个节点
elif cpn_obj.component_name.lower() in ("iteration", "loop"):
    _append_path(cpn_obj.get_start())

# 退出循环 → 到循环体的下游
elif cpn_obj.component_name.lower() == "exitloop":
    _extend_path(self.get_component(cpn["parent_id"])["downstream"])

# 有父组件的叶子节点 → 回到父组件的 start
elif not cpn["downstream"] and cpn_obj.get_parent():
    _append_path(cpn_obj.get_parent().get_start())

# 默认 → 推进到 downstream 列表
else:
    _extend_path(cpn["downstream"])

if cpn_obj.error():
    ex = cpn_obj.exception_handler()
    if ex and ex["goto"]:
        self.path.extend(ex["goto"])       # 跳转到指定的异常处理路径
        other_branch = True
    elif ex and ex["default_value"]:
        yield decorate("message", {"content": ex["default_value"]})
        # 输出默认值并继续
    else:
        self.error = cpn_obj.error()       # 标记全局错误，终止执行

def tts(self, tts_mdl, text):
    def clean_tts_text(text: str) -> str:
        # 清理控制字符、emoji、多余空格
        # 截断到 500 字符
    if not tts_mdl or not text:
        return None
    text = clean_tts_text(text)
    return synthesize_with_cache(tts_mdl, text)

async def get_files_async(self, files, layout_recognize=None) -> list[str]:
    # 图片文件 → image_to_base64 (data:image/...;base64,...)
    # 文档文件 → FileService.parse (文本提取)
    tasks = []
    for file in files:
        if file["mime_type"].find("image") >= 0:
            tasks.append(loop.run_in_executor(..., image_to_base64, file))
        else:
            tasks.append(loop.run_in_executor(..., parse_file, file))
    return await asyncio.gather(*tasks)

def tool_use_callback(self, agent_id, func_name, params, result, elapsed_time=None):
    # 记录每个组件调用工具的情况到 Redis
    # key: {task_id}-{message_id}-logs, TTL: 10分钟
    # 结构: [{"component_id": ..., "trace": [{path, tool_name, arguments, result, elapsed_time}]}]

def add_reference(self, chunks, doc_infos):
    # 去重后添加到 retrieval 栈的最顶层
    r = self.retrieval[-1]
    for ck in chunks_format({"chunks": chunks}):
        cid = hash_str2int(ck["id"], 500)
        if cid not in r:
            r["chunks"][cid] = ck
    for doc in doc_infos:
        if doc["doc_name"] not in r:
            r["doc_aggs"][doc["doc_name"]] = doc

def get_reference(self):
    return self.retrieval[-1] if self.retrieval else {"chunks": {}, "doc_aggs": {}}

def get_history(self, window_size):
    # 返回最近 window_size 轮的对话
    # window_size <= 0 返回空列表

def add_user_input(self, question):
    self.history.append(("user", question))
    self.globals["sys.history"].append(f"{self.history[-1][0]}: {self.history[-1][1]}")

def get_input(self, key=None):
    for var, o in self.get_input_elements().items():
        v = self.get_param(var)
        if isinstance(v, str) and self._canvas.is_reff(v):
            self.set_input_value(var, self._canvas.get_variable_value(v))
        elif isinstance(v, str) and re.search(self.variable_ref_patt, v):
            elements = self.get_input_elements_from_text(v)
            kv = {k: e.get('value', '') for k, e in elements.items()}
            self.set_input_value(var, self.string_format(v, kv))
        else:
            self.set_input_value(var, v)
    return res

COMPONENT_RENAMES = {
    "Splitter": "TokenChunker",
    "HierarchicalMerger": "TitleChunker",
    "PDFGenerator": "DocGenerator",
}
NODE_TYPE_RENAMES = {
    "splitterNode": "chunkerNode",
}

# 组件设置输出为一个 partial（包装了生成器/迭代器）
cpn_obj.set_output("content", partial(iter_chunks, generator))

# Canvas 检测到 partial，逐 chunk 消费
if isinstance(cpn_obj.output("content"), partial):
    stream = cpn_obj.output("content")()
    for m in stream:
        yield decorate("message", {"content": m})
    cpn_obj.set_output("content", _m)  # 最终替换为完整字符串

# 来自 deepdoc/parser/pdf_parser.py
sections = [(b["text"], f"@@{page}\t{left}\t{right}\t{top}\t{bottom}##") for b in boxes]

# 来自 rag/app/naive.py::Docx.__call__
sections = [(line.get("text"), line.get("image"), line.get("table")) for line in lines]

def naive_merge(sections, chunk_token_num=128, delimiter="\n。；！？", overlapped_percent=0):

def add_chunk(t, pos):
    tnum = num_tokens_from_string(t)
    # 如果文本太短（< 8 tokens），丢弃位置信息，避免碎片化
    if tnum < 8:
        pos = ""

    # 关键判断：当前 chunk 的 token 数是否超过阈值
    # 阈值 = chunk_token_num * (100 - overlapped_percent) / 100
    if cks[-1] == "" or tk_nums[-1] > chunk_token_num * (100 - overlapped_percent) / 100.:
        # 起新 chunk，先带上旧 chunk 的 overlap 尾巴
        if cks:
            overlapped = remove_tag(cks[-1])
            t = overlapped[int(len(overlapped) * (100 - overlapped_percent) / 100.):] + t
        if t.find(pos) < 0:
            t += pos
        cks.append(t)           # 创建新 chunk
        tk_nums.append(tnum)
    else:
        # 追加到当前 chunk（继续累积）
        if cks[-1].find(pos) < 0:
            t += pos
        cks[-1] += t            # 累积文本
        tk_nums[-1] += tnum

overlapped = remove_tag(cks[-1])         # 去除 PDF 位置标签，得到纯文本
# 取后半部分作为 overlap
t = overlapped[int(len(overlapped) * 0.8):] + t

custom_pattern = "##|---"    # 从 `##` 和 `---` 中提取

for sec, pos in sections:
    split_sec = re.split(r"(%s)" % custom_pattern, sec)
    for sub_sec in split_sec:
        cks.append("\n" + sub_sec)    # 每个分隔段直接成为独立 chunk

sections = [(text, image, table), ...]
     │
     ▼
cks = [
    {"text": "...", "ck_type": "text",   "tk_nums": 50},
    {"text": "...", "ck_type": "image",  "tk_nums": 0},      ← 图片不参与 token 合并
    {"text": "...", "ck_type": "table",  "tk_nums": 120},
    {"text": "...", "ck_type": "text",   "tk_nums": 80},
    ...
]

def _add_context(cks, idx, context_size):
    # 向前收集文本
    while prev >= 0 and remain_above > 0:
        if cks[prev]["ck_type"] == "text":
            piece = take_sentences_from_end(cks[prev]["text"], remain_above)
            parts_above.insert(0, piece)

    # 向后收集文本
    while after < len(cks) and remain_below > 0:
        if cks[after]["ck_type"] == "text":
            piece = take_sentences_from_start(cks[after]["text"], remain_below)
            parts_below.append(piece)

    cks[idx]["context_above"] = "".join(parts_above)
    cks[idx]["context_below"] = "".join(parts_below)

def _merge_cks(cks, chunk_token_num, has_custom):
    for i in range(len(cks)):
        ck_type = cks[i]["ck_type"]

        if ck_type == "text":
            # 上一个文本 chunk 还没满 → 合并
            if prev_text_ck >= 0 \
               and merged[prev_text_ck]["tk_nums"] < chunk_token_num \
               and not has_custom:         # 自定义分隔符模式下不合并
                merged[prev_text_ck]["text"] += cks[i]["text"]
                merged[prev_text_ck]["tk_nums"] += cks[i]["tk_nums"]
            else:
                merged.append(cks[i])      # 起新 chunk
        else:
            merged.append(cks[i])          # 图片/表格独立保留，不合并

self.delimiter_mode = "token_size"  # 默认：按 token 数累积合并
                     = "delimiter"  # 按分隔符切分
                     = "one"        # 一整块，不切

# 第一次切分：按主要分隔符（如 `\n`）
chunks = _split_text_by_pattern(payload, delimiter_pattern)

# 第二次切分：用 children_delimiters 对每个 chunk 再做切分
if custom_pattern:
    for c in cks:
        for text in _split_text_by_pattern(c, custom_pattern):
            docs.append({"text": text, "mom": c})
            # mom 字段保留父 chunk 原文，提供更大范围的上下文

def _attach_context_to_media_chunks(chunks, table_context_size, image_context_size):
    for i, chunk in enumerate(chunks):
        if chunk["ck_type"] not in {"table", "image"}:
            continue

        context_size = image_context_size if chunk["ck_type"] == "image" else table_context_size

        # 向前收集文本 chunks，直到满足 token 预算
        prev = i - 1
        while prev >= 0 and remain_above > 0:
            if prev_chunk["ck_type"] == "text":
                if prev_chunk["tk_nums"] >= remain_above:
                    # 只取末尾部分句子
                    parts_above.insert(0, _take_sentences(prev_chunk["text"], remain_above, from_end=True))
                else:
                    parts_above.insert(0, prev_chunk["text"])

        # 向后收集，对称逻辑
        after = i + 1
        while after < len(chunks) and remain_below > 0:
            # ... 对称逻辑

        chunk["context_above"] = "".join(parts_above)
        chunk["context_below"] = "".join(parts_below)

def tokenize(d, txt, eng):
    d["content_with_weight"] = txt
    # 去掉 HTML 表格标签后做粗粒度分词
    t = re.sub(r"</?(table|td|caption|tr|th)( [^<>]{0,12})?>", " ", txt)
    d["content_ltks"] = rag_tokenizer.tokenize(t)
    # 细粒度分词（用于模糊匹配）
    d["content_sm_ltks"] = rag_tokenizer.fine_grained_tokenize(d["content_ltks"])