from __future__ import annotations

# P0 chat runtime service.
# Dependencies are injected once from ui_server during app assembly so the
# large streaming/runtime logic is no longer owned by the entry module.

from services.context_compressor import (
    build_request_spec_from_messages,
    extract_request_spec,
    merge_request_spec,
    render_request_spec,
)

def _inject(context):
    globals().update(context)

async def stream_chat(tid,prompt):
    global last_llm_error
    usage_user_id=thread_user_ids.get(tid,"")

    pending_uploads=thread_upload_ids.pop(tid,[])
    upload_context=""

    if usage_user_id and pending_uploads:
        upload_context=await build_user_upload_context(
            usage_user_id,
            pending_uploads,
        )

        if upload_context:
            asyncio.create_task(
                safe_memory_event(
                    usage_user_id,
                    "attachment_used",
                    {
                        "thread_id":tid,
                        "upload_ids":pending_uploads[:10],
                        "count":len(pending_uploads[:10]),
                    },
                )
            )
    processing_context=""

    if (
        usage_user_id
        and pending_uploads
        and _quality_check_requested(prompt)
    ):
        task_id="proc_"+secrets.token_hex(8)

        await safe_memory_event(
            usage_user_id,
            "processing_started",
            {
                "task_id":task_id,
                "thread_id":tid,
                "operation":"quality_check",
                "upload_ids":
                    pending_uploads[:10],
                "status":"running",
            },
        )

        yield out(
            "status",
            {
                "text":
                    "正在使用本地 Python 检查上传数据…"
            },
        )

        try:
            processing_result=await asyncio.to_thread(
                _run_upload_quality_checks,
                usage_user_id,
                pending_uploads,
            )

            record=_compact_processing_record(
                processing_result
            )

            await safe_memory_event(
                usage_user_id,
                "processing_completed",
                {
                    "task_id":task_id,
                    "thread_id":tid,
                    "operation":
                        "quality_check",
                    "upload_ids":
                        pending_uploads[:10],
                    "status":"completed",
                    "result":record,
                },
            )

            processing_context=(
                "[USER_DATA_PROCESSING_RESULT]\n"
                "The following result was computed "
                "locally with Python from the user's "
                "uploaded file. Treat these computed "
                "values as authoritative for this "
                "answer. Do not estimate them from "
                "the raw file. Explain the result "
                "clearly in Chinese.\n"
                + json.dumps(
                    record,
                    ensure_ascii=False,
                    default=str,
                )
                + "\n[/USER_DATA_PROCESSING_RESULT]"
            )

            yield out(
                "status",
                {
                    "text":
                        "数据质检完成，正在整理结果…"
                },
            )

        except Exception as exc:
            await safe_memory_event(
                usage_user_id,
                "processing_failed",
                {
                    "task_id":task_id,
                    "thread_id":tid,
                    "operation":
                        "quality_check",
                    "upload_ids":
                        pending_uploads[:10],
                    "status":"failed",
                    "error":
                        str(exc)[:500],
                },
            )

            raise

    use_fisheries = _is_fisheries_prompt(prompt)
    use_ocean = _needs_ocean_mcp(prompt)
    mcp_error = None
    turn_error = ""
    export_tool_completed=False
    export_tool_results=[]
    try:
        hf_task = None
        if use_fisheries:
            yield out("status",{"text":"正在读取 Hugging Face 渔业数据…","stage":"source","step":2,"total_steps":5})
            hf_task = asyncio.create_task(build_hf_fisheries_context(prompt))

        # Marine MCP is loaded by CodeWhale from DEEPSEEK_MCP_CONFIG when the
        # runtime process starts.  Do NOT bootstrap it by asking the model to
        # call start_mcp_server inside every user thread: that creates a long
        # blocking turn and can make the browser/HF proxy drop the SSE stream.
        # A normal Ocean request goes straight to the real user turn; if the
        # model needs Ocean data it can call the already-registered mcp_marine_*
        # tools directly.
        if use_ocean:
            yield out("status",{"text":"Ocean 数据工具已就绪，正在处理请求…","stage":"source","step":2,"total_steps":5})

        grounded_prompt = prompt

        if processing_context:
            grounded_prompt += (
                "\n\n" + processing_context
            )
        elif upload_context:
            grounded_prompt += (
                "\n\n" + upload_context
            )
        if hf_task is not None:
            try:
                hf_context = await hf_task
                grounded_prompt = grounded_prompt + "\n\n" + hf_context

                if usage_user_id:
                    asyncio.create_task(
                        safe_memory_event(
                            usage_user_id,
                            "fisheries_query",
                            {
                                "thread_id":tid,
                                "source":"huggingface-fisheries",
                                "status":"completed",
                                "prompt_chars":len(prompt.strip()),
                            },
                        )
                    )
            except Exception as exc:
                grounded_prompt = (
                    prompt
                    + "\n\n[HF_FISHERIES_LIVE_CONTEXT_ERROR]\n"
                    + str(exc)
                    + "\n[/HF_FISHERIES_LIVE_CONTEXT_ERROR]"
                )

        try:
            det=await rjson(f"/v1/threads/{tid}")
            since=int(det.get("latest_seq") or 0)
            tr=await rjson(f"/v1/threads/{tid}/turns",method="POST",body={
                "prompt":grounded_prompt,
                "input_summary":prompt[:200],
                "model":MODEL,
                "mode":"agent",
                "allow_shell":False,
                "trust_mode":False,
                "auto_approve":False,
            })
        except Exception as exc:
            if not _is_codewhale_missing_thread_error(exc):
                raise
            history_text=_format_recent_thread_history(tid)
            if history_text:
                grounded_prompt=(
                    "[PRIOR_CONVERSATION_HISTORY]\n"
                    + history_text
                    + "\n[/PRIOR_CONVERSATION_HISTORY]\n\n"
                    + "[CONVERSATION_RECOVERY_INSTRUCTION]\n"
                    + "用户在本次对话中已经描述过下载/查询要求。上面是同一用户的完整对话记录；"
                    + "请把最新用户消息当作该对话的后续补充或修改继续处理，"
                    + "不要要求用户重新提交完整需求。\n"
                    + "[/CONVERSATION_RECOVERY_INSTRUCTION]\n\n"
                    + grounded_prompt
                )
            old_tid=tid
            tid=await recover_codewhale_thread(old_tid)
            log.warning(
                "CodeWhale thread expired and was rebuilt: "
                "old=%s new=%s",
                old_tid,
                tid,
            )
            yield out(
                "status",
                {
                    "text":
                        "检测到原对话会话已过期，正在自动保留上下文并恢复…",
                    "stage":"source",
                    "step":1,
                    "total_steps":5,
                },
            )
            yield out(
                "thread_recovered",
                {
                    "thread_id":tid,
                    "old_thread_id":old_tid,
                },
            )
            det=await rjson(f"/v1/threads/{tid}")
            since=int(det.get("latest_seq") or 0)
            tr=await rjson(f"/v1/threads/{tid}/turns",method="POST",body={
                "prompt":grounded_prompt,
                "input_summary":prompt[:200],
                "model":MODEL,
                "mode":"agent",
                "allow_shell":False,
                "trust_mode":False,
                "auto_approve":False,
            })
        turn=((tr or {}).get("turn") or {}).get("id")
        answer=""
        yield out("status",{"text":"Codex 正在处理…","stage":"compose","step":4,"total_steps":5})

        async for rec in events(tid,since):
            if turn and rec.get("turn_id") and rec["turn_id"]!=turn:
                continue

            e=rec.get("event")
            p=pl(rec)

            if e=="item.started":
                tool=_tool_name_from_payload(p)

                if tool.startswith("mcp_marine_"):
                    yield out("status",{"text":"正在查询学校 Ocean 数据服务器…","stage":"query","step":3,"total_steps":5})

                    if usage_user_id:
                        asyncio.create_task(
                            safe_memory_event(
                                usage_user_id,
                                _marine_event_type(tool),
                                {
                                    "thread_id":tid,
                                    "tool":tool,
                                    "operation":_marine_event_type(tool),
                                    "source":_data_source_from_prompt(prompt),
                                    "status":"started",
                                    "prompt_chars":len(prompt.strip()),
                                },
                            )
                        )

            if e=="item.completed":
                item=p.get("item") or {}
                summary=str(item.get("summary") or "")
                completed_tool=_tool_name_from_payload(p)
                try:
                    completed_raw=json.dumps(p,ensure_ascii=False,default=str)
                except Exception:
                    completed_raw=str(p)

                if (
                    completed_tool.endswith("_export")
                    or "mcp_marine_marine_export" in completed_raw
                ):
                    export_tool_completed=True
                    export_tool_results.append(completed_raw)

                if "MCP server 'marine' connected" in summary or "mcp_marine_" in summary:
                    marine_threads.add(tid)
                # Some Runtime versions can complete an agent_message item without
                # delivering a delta to this bridge. Recover the materialized final text.
                recovered=_agent_text_from_payload(p)
                if recovered and not answer:
                    answer=recovered

            if e=="item.delta" and p.get("kind")=="agent_message":
                d=p.get("delta") or ""
                if d:
                    answer+=d

            if e in {"item.failed","item.interrupted"}:
                err=_error_from_payload(p)
                if err:
                    turn_error=err

            if e=="approval.required":
                aid=p.get("approval_id") or p.get("id")
                tool=(
                    p.get("tool_name")
                    or ((p.get("tool") or {}).get("name") if isinstance(p.get("tool"),dict) else p.get("tool"))
                    or ((p.get("item") or {}).get("tool_name") if isinstance(p.get("item"),dict) else "")
                    or ""
                )
                if aid:
                    if tool.startswith("mcp_marine_"):
                        await approve(aid,"allow")
                    elif tool=="start_mcp_server":
                        # Static MCP config is authoritative.  Never start another
                        # MCP server dynamically inside an end-user thread.
                        await approve(aid,"deny")
                        yield out("status",{"text":"已阻止重复启动 Ocean MCP"})
                    else:
                        await approve(aid,"deny")
                        yield out("status",{"text":"已保持安全数据访问模式"})

            if e=="turn.lifecycle":
                st=_turn_status(p)
                if st in {"failed","canceled","interrupted"}:
                    detail=_error_from_payload(p) or turn_error or f"Turn {st}"
                    raise RuntimeError(detail)

            if e=="turn.completed":
                st=_turn_status(p)
                if st in {"failed","canceled","interrupted"}:
                    detail=_error_from_payload(p) or turn_error or f"Turn {st}"
                    raise RuntimeError(detail)
                if not answer.strip():
                    detail=_error_from_payload(p) or turn_error
                    if detail:
                        raise RuntimeError(detail)
                    raise RuntimeError(
                        "Codex 回合已结束，但 CodeWhale 没有产生 assistant 文本。"
                    )
                final_answer=_sanitize_final_answer(answer)
                if not final_answer:
                    raise RuntimeError("模型返回内容在输出清理后为空。")

                export_error=_ocean_export_execution_error(
                    prompt,
                    export_tool_completed=export_tool_completed,
                    tool_result_text="\n".join(export_tool_results),
                    final_answer=final_answer,
                )
                if export_error:
                    raise RuntimeError(export_error)

                if usage_user_id:
                    asyncio.create_task(
                        record_generated_download_assets(
                            usage_user_id,
                            tid,
                            prompt,
                            final_answer,
                        )
                    )

                last_llm_error=None
                yield out("token",{"text":final_answer})
                yield out("done",{"text":final_answer})
                return

        raise RuntimeError(turn_error or "Runtime stream ended early")
    except Exception as exc:
        last_llm_error=str(exc)
        log.exception(
            "chat failed: thread=%s model=%s ocean=%s fisheries=%s",
            tid, MODEL, use_ocean, use_fisheries,
        )
        yield out("error",{"text":_public_error(str(exc)),"stage":"chat"})


async def harness_stream_chat(tid,prompt):
    global last_llm_error

    usage_user_id=thread_user_ids.get(tid,"")
    pending_uploads=thread_upload_ids.pop(tid,[])

    upload_context=""
    processing_context=""

    if usage_user_id and pending_uploads:
        upload_context=await build_user_upload_context(
            usage_user_id,
            pending_uploads,
        )

        asyncio.create_task(
            safe_memory_event(
                usage_user_id,
                "attachment_used",
                {
                    "thread_id":tid,
                    "upload_ids":pending_uploads[:10],
                    "count":len(pending_uploads[:10]),
                    "runtime":"codex-harness",
                },
            )
        )

    if (
        usage_user_id
        and pending_uploads
        and _quality_check_requested(prompt)
    ):
        task_id="proc_"+secrets.token_hex(8)

        await safe_memory_event(
            usage_user_id,
            "processing_started",
            {
                "task_id":task_id,
                "thread_id":tid,
                "operation":"quality_check",
                "upload_ids":pending_uploads[:10],
                "status":"running",
                "runtime":"codex-harness",
            },
        )

        yield out(
            "status",
            {
                "text":
                    "正在使用本地 Python 检查数据…"
            },
        )

        try:
            processing_result=await asyncio.to_thread(
                _run_upload_quality_checks,
                usage_user_id,
                pending_uploads,
            )

            record=_compact_processing_record(
                processing_result
            )

            await safe_memory_event(
                usage_user_id,
                "processing_completed",
                {
                    "task_id":task_id,
                    "thread_id":tid,
                    "operation":"quality_check",
                    "upload_ids":pending_uploads[:10],
                    "status":"completed",
                    "runtime":"codex-harness",
                    "result":record,
                },
            )

            processing_context=(
                "[USER_DATA_PROCESSING_RESULT]\n"
                "These values were computed locally "
                "with Python from the current user's "
                "uploaded file. Use them as the "
                "authoritative result. Do not guess "
                "or recompute them mentally. Explain "
                "the result clearly in Chinese.\n"
                + json.dumps(
                    record,
                    ensure_ascii=False,
                    default=str,
                )
                + "\n[/USER_DATA_PROCESSING_RESULT]"
            )

            yield out(
                "status",
                {
                    "text":
                        "数据计算完成，Codex Harness 正在整理结果…"
                },
            )

        except Exception as exc:
            await safe_memory_event(
                usage_user_id,
                "processing_failed",
                {
                    "task_id":task_id,
                    "thread_id":tid,
                    "operation":"quality_check",
                    "upload_ids":pending_uploads[:10],
                    "status":"failed",
                    "runtime":"codex-harness",
                    "error":str(exc)[:500],
                },
            )
            raise

    app_context=thread_system_prompts.get(
        tid,
        USER_SYSTEM,
    )

    harness_prompt=(
        "[APPLICATION_CONTEXT]\n"
        + app_context
        + "\n[/APPLICATION_CONTEXT]\n\n"
        + "[CURRENT_USER_MESSAGE]\n"
        + prompt
        + "\n[/CURRENT_USER_MESSAGE]"
    )

    if processing_context:
        harness_prompt += (
            "\n\n" + processing_context
        )
    elif upload_context:
        harness_prompt += (
            "\n\n" + upload_context
        )

    yield out(
        "status",
        {
            "text":
                f"Codex Harness · {HARNESS_MODEL} 正在处理…"
        },
    )

    try:
        task=asyncio.create_task(
            asyncio.to_thread(
                dsh.run,
                harness_prompt,
                session_id=tid,
            )
        )

        while not task.done():
            try:
                await asyncio.wait_for(
                    asyncio.shield(task),
                    timeout=8,
                )
            except asyncio.TimeoutError:
                yield ": keepalive\n\n"

        result=await task

        final_answer=_sanitize_final_answer(
            result.final_response
        )

        if not final_answer:
            raise RuntimeError(
                "Codex Harness 没有返回有效文本。"
            )

        log.info(
            "Codex Harness completed: "
            "thread=%s model=%s reason=%s uploads=%s",
            tid,
            HARNESS_MODEL,
            result.finish_reason,
            len(pending_uploads),
        )

        last_llm_error=None

        yield out(
            "token",
            {"text":final_answer},
        )

        yield out(
            "done",
            {
                "text":final_answer,
                "runtime":"codex-harness",
                "model":HARNESS_MODEL,
                "finish_reason":result.finish_reason,
            },
        )

    except Exception as exc:
        last_llm_error=str(exc)

        log.exception(
            "Codex Harness failed: "
            "thread=%s model=%s uploads=%s",
            tid,
            HARNESS_MODEL,
            len(pending_uploads),
        )

        yield out(
            "error",
            {
                "text":
                    "Codex Harness 调用失败："
                    + str(exc)[:500],
                "stage":"harness",
            },
        )


DATA_REQUEST_TTL_SECONDS = 6 * 3600


def _context_spec_has_data(spec) -> bool:
    if not isinstance(spec, dict):
        return False
    return any(
        spec.get(key)
        for key in (
            "domain", "dataset", "variables", "region",
            "depth", "date", "format",
        )
    )


async def dispatch_chat_stream(tid,prompt):
    now=time.time()
    pending=thread_last_data_requests.get(tid) or {}
    if pending and now-float(pending.get("ts") or 0)>DATA_REQUEST_TTL_SECONDS:
        thread_last_data_requests.pop(tid, None)
        pending={}

    history_spec=build_request_spec_from_messages(
        thread_recent_history.get(tid) or []
    )
    context_spec=thread_compressed_contexts.get(tid) or {}
    if not _context_spec_has_data(context_spec):
        context_spec=history_spec
    if not _context_spec_has_data(context_spec) and pending.get("spec"):
        context_spec=pending.get("spec")

    explicit_data=(
        _needs_ocean_mcp(prompt)
        or _is_fisheries_prompt(prompt)
    )
    is_confirmation=_is_confirmation_prompt(prompt)
    is_amendment=bool(
        (not is_confirmation)
        and pending.get("prompt")
        and _is_data_request_followup(prompt)
    )
    resumed=bool(
        pending.get("prompt")
        and (is_confirmation or is_amendment)
    )

    # Remember every data request and each short amendment that depends on it.
    # This lets follow-ups such as “补充区域…”/“把格式改成 CSV” stay in the
    # same data runtime and be merged with the original parameters instead of
    # being treated as brand-new unrelated messages.
    if is_amendment:
        merged_spec=merge_request_spec(
            context_spec,
            prompt,
            revision=True,
        )
        thread_compressed_contexts[tid]=merged_spec
        thread_last_data_requests[tid]={
            "prompt":render_request_spec(merged_spec),
            "ts":now,
            "spec":merged_spec,
        }
    elif explicit_data and not resumed:
        parsed_spec=extract_request_spec(prompt)
        if _context_spec_has_data(parsed_spec):
            merged_spec=merge_request_spec(
                None,
                prompt,
                revision=False,
            )
            thread_compressed_contexts[tid]=merged_spec
        else:
            merged_spec={"raw_request":str(prompt or "")[:1600]}
            thread_compressed_contexts[tid]=merged_spec
        thread_last_data_requests[tid]={
            "prompt":render_request_spec(merged_spec) or prompt,
            "ts":now,
            "spec":merged_spec,
        }

    # Fast-path multi-day Ocean exports into a persistent batch job instead of
    # keeping one Agent turn open for dozens of serial tool calls.
    batch_spec=parse_ocean_batch_request(prompt)
    if batch_spec:
        if not resumed:
            batch_spec_context=thread_compressed_contexts.get(tid) or {}
            thread_last_data_requests[tid]={
                "prompt":(
                    render_request_spec(batch_spec_context)
                    or prompt
                ),
                "ts":now,
                "spec":batch_spec_context,
            }
        uid=thread_user_ids.get(tid,"")
        job=await OCEAN_BATCH_MANAGER.create(uid,prompt,batch_spec,tid)
        yield out("status",{
            "text":f"已拆分为 {job['total']} 个单日子任务，后台并行执行中…",
            "stage":"query","step":3,"total_steps":5,
        })
        yield out("batch_job",job)
        variables="、".join(batch_spec.get("variables") or [])
        bbox=batch_spec.get("bbox") or {}
        region_note=(
            f"区域：{batch_spec.get('region_name')} "
            f"({bbox.get('lon_min')}°E–{bbox.get('lon_max')}°E, "
            f"{bbox.get('lat_min')}°N–{bbox.get('lat_max')}°N)。"
            if batch_spec.get("region_name") else ""
        )
        text=(
            f"已创建 Ocean 批量导出任务，共 {job['total']} 个单日子任务。"
            f"范围 {batch_spec['start_date']} 至 {batch_spec['end_date']}，变量：{variables}。"
            + region_note +
            "学校 Ocean 接口只接收单日 YYYY-MM-DD；平台已自动把月份/日期范围拆成单日并行执行，"
            "不会再要求您补充某一天。任务卡会实时显示完成、失败、执行中和等待数量。"
        )
        yield out("done",{"text":text,"batch_job_id":job["job_id"]})
        return

    has_upload = bool(thread_upload_ids.get(tid))
    compact_spec_text=render_request_spec(
        thread_compressed_contexts.get(tid) or context_spec
    )
    routed_prompt=prompt
    if is_confirmation and resumed:
        thread_last_data_requests[tid]={
            "prompt":compact_spec_text or str(pending["prompt"]),
            "ts":now,
            "spec":thread_compressed_contexts.get(tid) or context_spec,
        }
        routed_prompt=(
            "请下载或继续执行以下已确认的数据请求，保留全部参数：\n\n"
            + (compact_spec_text or str(pending["prompt"]))
            + "\n\n[USER_CONFIRMATION]\n"
            + "用户刚刚回复确认。请立即继续执行上一项数据查询或导出，"
            + "沿用已经给出的日期、区域、变量和数据源，不要再次询问确认。\n"
            + "[/USER_CONFIRMATION]"
        )
    elif is_amendment:
        routed_prompt=(
            "请下载或继续执行以下数据请求，保留全部参数：\n\n"
            + str(
                thread_last_data_requests[tid].get("prompt")
                or pending["prompt"]
                or ""
            )
            + "\n\n[AMENDMENT_DIRECTIVE]\n"
            + "用户正在补充或修改上一项数据下载/查询请求。"
            + "请把“最近补充/修改”作为本次唯一变更，其余条件继续沿用卡片中的参数；"
            + "按最新要求立即执行，不要要求用户重新提交完整需求。\n"
            + "[/AMENDMENT_DIRECTIVE]"
        )
    routed_prompt=_apply_ocean_export_defaults(routed_prompt)

    keep_codewhale = (
        dsh is None
        or (
            not has_upload
            and (
                _needs_ocean_mcp(routed_prompt)
                or _is_fisheries_prompt(routed_prompt)
                or resumed
            )
        )
    )

    if keep_codewhale:
        async for chunk in stream_chat(tid, routed_prompt):
            yield chunk
        return

    async for chunk in harness_stream_chat(tid, routed_prompt):
        yield chunk


def init_chat_runtime(context):
    _inject(context)
    return stream_chat, harness_stream_chat, dispatch_chat_stream