"""Agentic loop for the Code tab — Claude-CLI style.

Drives the provider with the file/command tools, routing every write/run action
through the permission gate. Read-only tools run without prompting. For models
without native tool-calling, a text fallback protocol is supported: the model
emits a line ``@@TOOL <name> <json-args>``.
"""
from __future__ import annotations

import json
import re
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional

from ..providers.base import Provider
from . import agent_roles
from . import agent_security
from .ms365_tools import MS365_WRITE_TOOLS
from .permissions import PermissionGate
from .plan import UPDATE_PLAN_SPEC, normalize_plan_steps
from .tools import TOOL_SPECS, WRITE_TOOLS, ToolContext, describe_action, execute_tool

EmitFn = Callable[[Dict[str, Any]], None]
CancelFn = Callable[[], bool]

MAX_STEPS = 40  # headroom for diagnose → fix → retry loops
_TOOL_LINE = re.compile(r"@@TOOL\s+(\w+)\s+(\{.*\})", re.DOTALL)


def code_system_prompt(workdir: Path, has_memory: bool = False, plan: bool = False,
                       has_plan_tool: bool = False, has_ms365: bool = False) -> str:
    names = ", ".join(t.name for t in TOOL_SPECS)
    plan_note = ("PLAN MODE: only analyze and propose a detailed plan; do NOT write files or run "
                 "commands. When the user asks to gencode/implement, the app switches to ACT.\n"
                 if plan else "")
    memory_note = ""
    if has_memory:
        memory_note = (
            "You also have 'codebase memory' (cmem_*): cmem_get_architecture, cmem_search_graph, "
            "cmem_trace_path, cmem_get_code_snippet, cmem_search_code, cmem_query_graph. "
            "Prefer these to understand code structure (callers/callees, where functions live) "
            "instead of reading/grepping file by file — faster and fewer tokens.\n"
        )
    plan_tool_note = ""
    if has_plan_tool:
        plan_tool_note = (
            "A step checklist is shown to the user. Call update_plan(steps=[{title, status}]) to "
            "keep it in sync as you work — mark the current step 'running', then 'done' when "
            "finished (status is one of pending/running/done).\n"
        )
    ms365_note = ""
    if has_ms365:
        ms365_note = (
            "The user has signed in to Microsoft 365 and enabled some ms365__* tools (Outlook / "
            "Teams / OneDrive / SharePoint / meeting transcripts, via the built-in MS365 MCP "
            "server). Use them whenever the request involves that data — don't say you can't "
            "access it.\n"
        )
    return (
        "You are Cowork Code — a coding assistant that works like a CLI agent.\n"
        f"Current working folder: {workdir}\n"
        + plan_note +
        "You can call the tools: " + names + ".\n"
        + memory_note + plan_tool_note + ms365_note +
        "Break work into steps, read files before editing, and explain each step briefly.\n"
        "For changes to an existing file, prefer edit_file (replace an exact snippet) over "
        "rewriting the whole file with write_file; use write_file only for new files or full "
        "rewrites. Always read_file first so old_string matches exactly.\n"
        "If a task needs a Python library that isn't installed, install it yourself with the "
        "install_package tool (or `pip install` via run_command) and continue — never ask the "
        "user to install libraries by hand. Commands and installs run inside this project's own "
        "isolated '.venv' (created automatically on first use), separate from other projects.\n"
        "When you must GENERATE a deliverable (e.g. .pptx/.docx/.xlsx/.pdf/images) by writing and "
        "running a script: put the generator script and any temporary files in a '.scratch/' "
        "subfolder, run it so the final file lands in the working folder, then DELETE the "
        "'.scratch/' folder. Only the final requested file(s) should remain — never leave "
        "generator scripts or intermediate files behind.\n"
        "Every path must stay inside the working folder.\n"
        "If a command or tool fails, do NOT stop and hand the error back to the user — read the "
        "error, fix the cause (edit the code, install a missing package, correct the command) and "
        "retry. Keep iterating until the task actually works, then run it once more so you can "
        "show the real output the user asked for.\n"
        "If the environment can't call tools directly, emit exactly ONE line of the form:\n"
        "@@TOOL <tool_name> {\"param\": \"value\"}\n"
        "When the task is complete, reply to the user in plain text — include the produced output "
        "— and stop calling tools. Answer in the user's language."
    )


_SKILLS_TAG = "[[ACTIVE_SKILLS]]"
_RULES_TAG = "[[SECURITY_RULES]]"
_PROJECT_TAG = "[project-context]"


def _apply_skills(messages: List[Dict[str, Any]], skills_text: str) -> None:
    """Insert/refresh a single system message carrying the enabled skills."""
    messages[:] = [
        m for m in messages
        if not (m.get("role") == "system" and str(m.get("content", "")).startswith(_SKILLS_TAG))
    ]
    if not skills_text.strip():
        return
    block = {
        "role": "system",
        "content": f"{_SKILLS_TAG}\nThe user enabled the following skills — follow them:\n\n{skills_text}",
    }
    insert_at = 1 if messages and messages[0].get("role") == "system" else 0
    messages.insert(insert_at, block)


def _apply_security_rules(messages: List[Dict[str, Any]], rules_text: str) -> None:
    """Insert/refresh a single system message carrying the external security/
    restriction rules (core/security_rules.py) — same pattern as _apply_skills,
    but for mandatory guardrails rather than opt-in behaviors."""
    messages[:] = [
        m for m in messages
        if not (m.get("role") == "system" and str(m.get("content", "")).startswith(_RULES_TAG))
    ]
    if not rules_text.strip():
        return
    block = {
        "role": "system",
        "content": (f"{_RULES_TAG}\nMandatory security/restriction rules — check every request "
                    "and action against these BEFORE acting; refuse or ask for clarification "
                    f"instead of proceeding if something would violate them:\n\n{rules_text}"),
    }
    insert_at = 1 if messages and messages[0].get("role") == "system" else 0
    messages.insert(insert_at, block)


def _apply_project_context(messages: List[Dict[str, Any]], context_text: str) -> None:
    """Insert/refresh a single system message carrying the project's shared
    instructions (Claude-Projects style) — same pattern as _apply_skills, so
    editing the project context takes effect on the NEXT turn of every one of
    its threads, without duplicating blocks in long conversations."""
    messages[:] = [
        m for m in messages
        if not (m.get("role") == "system" and str(m.get("content", "")).startswith(_PROJECT_TAG))
    ]
    if not (context_text or "").strip():
        return
    block = {"role": "system", "content": f"{_PROJECT_TAG}\n{context_text.strip()}"}
    insert_at = 1 if messages and messages[0].get("role") == "system" else 0
    messages.insert(insert_at, block)


_RECOVERY_NOTE = (
    "[auto-recovery] The previous attempt hit an error. Review the conversation so "
    "far, redo the most recent step if it looks incomplete or wrong, and fix any "
    "issue before continuing."
)


def _call_provider_with_recovery(provider: Provider, messages: List[Dict[str, Any]],
                                 tools, on_text, cancel, on_reasoning=None,
                                 max_retries: int = 1) -> Dict[str, Any]:
    """``provider.chat(...)`` with ONE bounded, silent recovery attempt: if the
    call raises (a dropped connection, an exhausted rate-limit wait, a
    momentarily unreachable gateway, ...), retry once with a short recovery
    note appended ONLY to that retry's OWN copy of ``messages`` — the caller's
    ``messages`` list is never mutated, so the note never leaks into the
    real, persisted conversation. Exhausting the retry re-raises the original
    exception unchanged, preserving existing failure handling (e.g. the
    "model not found" restore-to-composer flow).

    ``on_reasoning`` is only forwarded when given — ``run_code`` never passed
    it before this helper existed, and some lightweight test doubles for
    ``Provider`` don't accept the keyword at all; omitting it when unset
    keeps every existing call site's exact prior calling convention."""
    kwargs = {"tools": tools, "on_text": on_text, "cancel": cancel}
    if on_reasoning is not None:
        kwargs["on_reasoning"] = on_reasoning
    try:
        return provider.chat(messages, **kwargs)
    except Exception:
        if max_retries <= 0 or cancel():
            raise
        recovery = messages + [{"role": "user", "content": _RECOVERY_NOTE}]
        return _call_provider_with_recovery(provider, recovery, tools, on_text, cancel,
                                            on_reasoning, max_retries - 1)


def _parse_react(content: str) -> List[Dict[str, Any]]:
    """Extract @@TOOL fallback calls from assistant text."""
    calls: List[Dict[str, Any]] = []
    for i, match in enumerate(_TOOL_LINE.finditer(content or "")):
        name = match.group(1)
        try:
            args = json.loads(match.group(2))
        except json.JSONDecodeError:
            continue
        calls.append({"id": f"react_{i}", "name": name, "arguments": args})
    return calls


def run_code(
    provider: Provider,
    messages: List[Dict[str, Any]],
    ctx: ToolContext,
    gate: PermissionGate,
    emit: EmitFn,
    cancel: Optional[CancelFn] = None,
    max_steps: int = MAX_STEPS,
    extra_tools: Optional[List] = None,
    extra_executor=None,
    skills_text: str = "",
    rules_text: str = "",
    project_context: str = "",
    plan: bool = False,
    security_config=None,
) -> List[Dict[str, Any]]:
    """``security_config`` is the app's ``AppConfig`` — enables the same
    active guardrails as ``chat_agent.run_cowork`` (see its docstring).
    ``None`` (the default) disables both checks."""
    cancel = cancel or (lambda: False)
    extra_tools = extra_tools or []
    extra_names = {t.name for t in extra_tools}
    # update_plan drives the Plan panel (same as run_cowork) — always
    # available, not opt-in via extra_tools, so unattended Schedule Tasks
    # running through the code agent also get the completion self-check.
    from .tools import enabled_tool_specs
    all_tools = enabled_tool_specs(security_config) + extra_tools + [UPDATE_PLAN_SPEC]
    # MS365 tools that send/write data (mail, Teams messages, OneDrive
    # writes) are gated exactly like write_file/run_command — only the
    # read/list ms365 tools count as "read-only, never confirm". Names are
    # the MCP-qualified "ms365__*" form the agent sees (see ms365_tools.py).
    gated_tools = WRITE_TOOLS | MS365_WRITE_TOOLS
    # In PLAN mode, don't advertise write/run tools (analysis only).
    advertised = [t for t in all_tools if t.name not in gated_tools] if plan else all_tools
    has_memory = any(t.name.startswith("cmem_") for t in extra_tools)
    has_plan_tool = True
    has_ms365 = any(t.name.startswith("ms365_") for t in extra_tools)
    if not messages or messages[0].get("role") != "system":
        messages.insert(0, {"role": "system",
                            "content": code_system_prompt(
                                ctx.workdir, has_memory, plan, has_plan_tool, has_ms365)})
    _apply_skills(messages, skills_text)
    # The CODE agent uses RULEforCode.md (empty for now), NOT Cowork's
    # RULEBASE.md — its safety comes from the sandbox until code rules exist.
    from .security_rules import load_code_rules
    _apply_security_rules(messages, rules_text or load_code_rules())
    _apply_project_context(messages, project_context)
    # Active guardrail — reviews the request itself and can refuse to proceed
    # at all. Raises SecurityBlocked on a violation. agent_kind="code" → the
    # code rulebase (RULEforCode.md), so RULEBASE.md's rules don't gate coding.
    agent_security.enforce_prompt(provider, messages, security_config, emit, agent_kind="code")

    for _ in range(max_steps):
        if cancel():
            break

        def on_text(piece: str) -> None:
            emit({"type": "text", "delta": piece})

        assistant = _call_provider_with_recovery(provider, messages, advertised, on_text, cancel)
        if not assistant.get("tool_calls"):
            fallback = _parse_react(assistant.get("content", ""))
            if fallback:
                assistant["tool_calls"] = fallback
        messages.append(assistant)
        if not assistant.get("tool_calls") and not (assistant.get("content") or "").strip():
            # Reasoning-only reply with no answer and no tool call — never leave a
            # blank final turn (a Schedule Task run reads this back as its final
            # answer, so a blank one silently produces "(no output)").
            assistant["content"] = "*(model returned only its reasoning — try rephrasing)*"
        emit({"type": "assistant_done", "content": assistant.get("content", "")})

        tool_calls = assistant.get("tool_calls") or []
        if not tool_calls:
            break

        for tc in tool_calls:
            if cancel():
                return messages
            name, args, tc_id = tc["name"], tc.get("arguments", {}), tc["id"]
            # update_plan drives the Plan panel only — no chat bubble, no file.
            if name == "update_plan":
                steps = normalize_plan_steps(args.get("steps"))
                emit({"type": "plan_set", "steps": steps})
                messages.append({"role": "tool", "tool_call_id": tc_id, "name": name,
                                 "content": "Plan updated."})
                continue
            if plan and name in gated_tools:
                emit({"type": "tool_result", "id": tc_id, "name": name, "ok": False,
                      "output": "PLAN mode: action skipped (switch to Act to execute)."})
                messages.append({"role": "tool", "tool_call_id": tc_id, "name": name,
                                 "content": "PLAN mode: not executed."})
                continue
            is_extra = name in extra_names
            preview = ({"kind": "info", "title": name, "text": str(args)}
                       if is_extra else describe_action(ctx, name, args))
            emit({"type": "tool_proposed", "id": tc_id, "name": name, "args": args, "preview": preview})
            # Active guardrail on run_command/install_package, checked BEFORE
            # asking the user to confirm — a no-op for every other tool or
            # when disabled. Raises SecurityBlocked on a violation. Code agent
            # → RULEforCode.md rulebase (not Cowork's RULEBASE.md).
            agent_security.enforce_command(provider, name, args, security_config, emit,
                                           agent_kind="code")

            if name in gated_tools:
                approved = gate.request({"id": tc_id, "name": name, "args": args, "preview": preview})
            else:
                approved = True  # read-only tools (incl. codebase memory) never confirm

            if cancel():
                return messages

            if not approved:
                result = {"ok": False, "output": "User rejected the action."}
            else:
                emit({"type": "tool_start", "id": tc_id, "name": name})
                if is_extra and extra_executor is not None:
                    result = extra_executor(name, args)
                else:
                    def on_output(line: str, _id=tc_id, _name=name) -> None:
                        emit({"type": "tool_output", "id": _id, "name": _name, "delta": line})
                    result = execute_tool(ctx, name, args, cancel=cancel, on_output=on_output,
                                          agent_role=agent_roles.CODE)

            evt = {
                "type": "tool_result", "id": tc_id, "name": name,
                "ok": result["ok"], "output": result["output"],
            }
            if isinstance(args, dict) and args.get("path"):
                # absolute path so the UI can open the containing folder
                evt["path"] = str(ctx.workdir / str(args["path"]))
            if isinstance(result, dict) and result.get("produced"):
                evt["produced"] = result["produced"]
            emit(evt)
            messages.append({
                "role": "tool", "tool_call_id": tc_id, "name": name,
                "content": result["output"],
            })
    return messages