# agent_research.py - Research flow: gate, trim, scratch path, token estimate. Used by aiagent/agent_loop.
from __future__ import annotations

import os
import re
import json
from datetime import datetime
from utils.transcript_roles import tool_result_text_for_trim, is_tool_result_carrier_message

# Legacy name — message bodies + assistant tool_calls only (see utils.context_token_budget).
from utils.context_token_budget import estimate_message_tokens as estimate_tokens

# Research report gate: used by run_raw_agent and by tests.
MIN_SCRATCH_APPENDS_FOR_REPORT = 2
MIN_SCRATCH_FILE_SIZE = 2000
# Extra deep_web_search / read_and_analyze_page calls allowed when at depth cap but scratch/report gate not satisfied.
RESEARCH_SOFT_DEPTH_OVERFLOW = 3


def research_harness_eligible(
    *,
    research_mode: bool,
    subagent_mode: bool,
    kanban_worker_mode: bool,
    task_metadata: dict | None = None,
) -> bool:
    """True when research depth caps, state tracking, and post-loop report apply."""
    if not research_mode:
        return False
    if not subagent_mode:
        return True
    if not kanban_worker_mode:
        return False
    if task_metadata is not None:
        try:
            from utils.council_orchestrator import is_council_task_metadata

            if is_council_task_metadata(task_metadata):
                return False
        except Exception:
            pass
        kind = str((task_metadata or {}).get("task_kind") or "").strip().lower()
        if kind and kind != "research":
            return False
        if kind == "research":
            assignee = str((task_metadata or {}).get("_assignee") or "").strip()
            if assignee and assignee != "research":
                return False
    return True


def council_researcher_tools_must_be_stripped(
    *,
    kanban_worker_mode: bool,
    task_metadata: dict | None,
    ledger_last_phase: str | None,
) -> bool:
    """True when Athena is in CLOSE — tools stay off until kanban-result (survives guard retries)."""
    if not kanban_worker_mode:
        return False
    try:
        from utils.council_orchestrator import is_council_task_metadata

        tm = task_metadata if isinstance(task_metadata, dict) else {}
        if not is_council_task_metadata(tm):
            return False
        if str(tm.get("council_seat") or "").strip().lower() != "researcher":
            return False
        return str(ledger_last_phase or "").strip().upper() == "CLOSE"
    except Exception:
        return False


def council_researcher_close_nudge_eligible(
    *,
    kanban_worker_mode: bool,
    task_metadata: dict | None,
    research_tools_used: bool,
    round_: int,
    min_round: int = 3,
    research_gather_tool_count: int = 0,
    search_budget_exhausted: bool = False,
) -> bool:
    """Inject a hard CLOSE nudge so Athena finishes with kanban-result, not more tools."""
    if not kanban_worker_mode or int(round_) < 1:
        return False
    try:
        from agent_loop_guardrails import council_researcher_gather_sufficient
        from utils.council_orchestrator import (
            council_researcher_has_preseed,
            is_council_task_metadata,
        )

        tm = task_metadata if isinstance(task_metadata, dict) else {}
        if not is_council_task_metadata(tm):
            return False
        if str(tm.get("council_seat") or "").strip().lower() != "researcher":
            return False
        has_preseed = council_researcher_has_preseed(tm)
        min_round_effective = 1 if has_preseed else int(min_round)
        if int(round_) < min_round_effective:
            return False
        if not research_tools_used and not has_preseed:
            return False
        return council_researcher_gather_sufficient(
            research_gather_tool_count,
            search_budget_exhausted=search_budget_exhausted,
            has_preseed=has_preseed,
        )
    except Exception:
        return False


def research_scratch_reminder_eligible(
    *,
    research_mode: bool,
    subagent_mode: bool,
    kanban_worker_mode: bool,
    task_metadata: dict | None = None,
    research_tools_used: bool = False,
    append_scratch_count: int = 0,
) -> bool:
    """Whether the harness should inject the append_research_scratch reminder.

    Kanban research workers are subagents; they must not use the legacy
    ``not subagent_mode`` gate (regression guard for completion-controller PR2).
    """
    if not research_tools_used or int(append_scratch_count) > 0:
        return False
    return research_harness_eligible(
        research_mode=research_mode,
        subagent_mode=subagent_mode,
        kanban_worker_mode=kanban_worker_mode,
        task_metadata=task_metadata,
    )


# Keep at most this many large tool results in full; older ones replaced with placeholder.
_LARGE_TOOL_RESULT_THRESHOLD = 12_000
_LARGE_TOOL_RESULT_KEEP_LAST_N = 1

# Research tools: keep more context so the agent can save before losing it.
_RESEARCH_TOOL_NAMES = frozenset(("deep_web_search", "read_and_analyze_page", "analyze_multiple_sources", "append_research_scratch"))
_RESEARCH_INDICATOR_TOOLS = _RESEARCH_TOOL_NAMES | frozenset(("search_web", "search_web_public", "read_webpage"))
_RESEARCH_DEPTH_TOOLS = frozenset(("deep_web_search", "read_and_analyze_page", "analyze_multiple_sources", "search_web", "search_web_public", "read_webpage"))
_RESEARCH_LARGE_RESULT_KEEP_LAST_N = 2
_RESEARCH_TRIM_MIN_CHARS = 600
_SAVE_RESEARCH_REMINDER_ROUND = 6
_CTX_SMALL_THRESHOLD_TOKENS = 32_000

_LARGE_RESEARCH_PLACEHOLDER_BASE = (
    "[Research result from an earlier round omitted to reduce request size. "
    "Do not re-search or re-read; use your accumulated findings. "
    "If you have not yet saved, call **save_research_findings** or **generate_final_research_report** now with what you have; then you can continue with more sources if needed. "
)

_SCRATCH_PATH_RE = re.compile(r"Scratch file:\s*(\S+)", re.IGNORECASE)


def _stable_placeholder_appendix(text: str) -> str:
    """Short stable reference line for trimmed tool results (paths, scratch file)."""
    if not text:
        return ""
    m = _SCRATCH_PATH_RE.search(text)
    if m:
        return f" ref_scratch_path={m.group(1).strip()[:200]}"
    lines = [ln.strip() for ln in text.splitlines() if ln.strip()]
    for ln in lines[:12]:
        if "saved to:" in ln.lower() or ln.startswith("/") or ln.startswith("~/"):
            return f" ref_line={ln[:180]}"
    return ""


def write_scratch_from_findings(topic: str, findings: str) -> str | None:
    """Write accumulated findings to a scratch file and return the path, or None on failure."""
    from utils.path_safety import get_user_output_dir, safe_resolve_path
    base = get_user_output_dir()
    safe_topic = re.sub(r"[^\w\s-]", "", (topic or "research").strip())[:60].replace(" ", "_")
    date = datetime.now().strftime("%Y-%m-%d")
    relative = f"research/_scratch/{safe_topic}_{date}.md"
    abs_path, err = safe_resolve_path(base, relative, allow_write=True, create_parents=True)
    if err or not abs_path:
        return None
    try:
        with open(abs_path, "w", encoding="utf-8") as f:
            f.write(findings)
        return abs_path
    except Exception:
        return None


def find_latest_scratch_path(messages: list) -> str | None:
    """Super-robust: scan ALL tool results for scratch path; multiple patterns; validate file exists."""
    for idx, msg in enumerate(reversed(messages)):
        content = tool_result_text_for_trim(msg)
        if not isinstance(content, str) or len(content) < 20:
            continue
        patterns = [
            r"Scratch file:\s*([^\s(]+)",
            r"scratch_path\s*[:=]\s*['\"]?([^'\"\s()]+)['\"]?",
            r"(/[^'\"]+?_scratch/[^'\"]+?\.md)",
        ]
        for pat in patterns:
            m = re.search(pat, content, re.IGNORECASE)
            if m:
                candidate = m.group(1).strip().rstrip(")")
                expanded = os.path.expanduser(candidate)
                if os.path.isfile(expanded):
                    return expanded
    return None


def _research_report_gate_block_reason(
    append_scratch_count: int,
    scratch_path_arg: str | None,
    messages: list,
    *,
    min_appends: int = MIN_SCRATCH_APPENDS_FOR_REPORT,
    min_file_size: int = MIN_SCRATCH_FILE_SIZE,
    waive_thin_scratch: bool = False,
) -> str | None:
    """
    Return None if generate_final_research_report is allowed; else return the block message.
    Used by run_raw_agent and by tests.

    When *waive_thin_scratch* is True (after depth-cap recovery), the minimum scratch *file size*
    check is skipped once append count is satisfied — findings must still be non-empty in the tool.
    """
    if append_scratch_count < min_appends:
        return (
            f"Cannot generate report yet — you must call append_research_scratch at least "
            f"{min_appends} times first (you have {append_scratch_count}). "
            f"Right now, call append_research_scratch(topic, content=<detailed synthesis of everything "
            f"you've found so far — full paragraphs with specific facts, numbers, comparisons>) "
            f"to save your findings to the scratch file. Do this {min_appends - append_scratch_count} "
            f"more time(s), then call generate_final_research_report with scratch_path."
        )
    if waive_thin_scratch:
        return None
    scratch_path = scratch_path_arg or find_latest_scratch_path(messages)
    if not scratch_path:
        return None
    try:
        expanded = os.path.expanduser(str(scratch_path))
        size = os.path.getsize(expanded) if os.path.isfile(expanded) else 0
    except Exception:
        size = 0
    if size < min_file_size:
        return (
            f"Scratch file is too thin ({size:,} chars, need at least {min_file_size:,}). "
            f"Do more research (search + read pages), then call append_research_scratch with a DETAILED "
            f"synthesis — full paragraphs with specific facts, numbers, model names, benchmarks. "
            f"After the file has enough content, call generate_final_research_report again."
        )
    return None


def trim_old_large_tool_results(
    messages: list,
    research_retain: list | None = None,
    research_keep_last_n: int | None = None,
    *,
    other_placeholder_preview_chars: int | None = None,
    research_placeholder_flat_chars: int | None = None,
    other_keep_subtract: int = 0,
    pinned_tool_call_ids: set[str] | None = None,
) -> list[tuple[int, int, int]]:
    """Replace old tool messages with very large content by short placeholders so the next API request doesn't hit 500.
    Research tool results keep the last N in full; other large results keep 1.
    If research_retain is provided, append full content of every research tool result (trimmed and kept) so force-save can use real content instead of placeholders.
    research_keep_last_n overrides _RESEARCH_LARGE_RESULT_KEEP_LAST_N when provided (for context-aware scaling).
    When ``other_placeholder_preview_chars`` / ``research_placeholder_flat_chars`` are None, defaults match legacy (120 / 300).
    ``other_keep_subtract``: subtract from the standard ``keep_last_n`` for non-research blobs (tier2 tightening).
    ``pinned_tool_call_ids``: native ``tool_call_id``s of handoff-evidence results to leave full
    (Phase 2b); they are skipped by both placeholder loops. The set is byte-bounded upstream so
    pinning stays 500-safe.
    Returns list of (index, old_content_len, new_content_len) for incremental token estimate."""
    deltas: list[tuple[int, int, int]] = []
    _pins = pinned_tool_call_ids or set()
    tool_call_id_to_name = {}
    for msg in messages:
        if not isinstance(msg, dict) or msg.get("role") != "assistant":
            continue
        for tc in msg.get("tool_calls") or []:
            tid = tc.get("id") if isinstance(tc, dict) else getattr(tc, "id", None)
            fn = tc.get("function") if isinstance(tc, dict) else getattr(tc, "function", None)
            name = fn.get("name") if isinstance(fn, dict) else (getattr(fn, "name", None) if fn else None)
            if tid and name:
                tool_call_id_to_name[tid] = name
    xml_result_index_to_name: dict[int, str] = {}
    for i, msg in enumerate(messages):
        if not is_tool_result_carrier_message(msg) or i <= 0:
            continue
        prev = messages[i - 1]
        if not isinstance(prev, dict) or (prev.get("role") or "").strip() != "assistant":
            continue
        calls = prev.get("tool_calls") or []
        if not calls:
            content = prev.get("content") or ""
            if isinstance(content, str) and "<tool_call" in content.lower():
                try:
                    from agent_loop import _parse_xml_tool_calls  # lazy import to avoid hard cycle at module import time
                    calls = _parse_xml_tool_calls(content)
                except Exception:
                    calls = []
        if len(calls) == 1 and isinstance(calls[0], dict):
            fn = calls[0].get("function") if isinstance(calls[0], dict) else None
            name = (
                (fn.get("name") if isinstance(fn, dict) else None)
                or calls[0].get("name")
                or ""
            )
            if name:
                xml_result_index_to_name[i] = str(name)

    research_indices = []
    other_indices = []
    for i, msg in enumerate(messages):
        if not isinstance(msg, dict):
            continue
        content = tool_result_text_for_trim(msg)
        is_str = isinstance(content, str)
        clen = len(content) if is_str else (len(content) if isinstance(content, (list, tuple)) else 0)
        tc_id = msg.get("tool_call_id") or ""
        name = tool_call_id_to_name.get(tc_id) or xml_result_index_to_name.get(i)
        passes_large = is_str and clen >= _LARGE_TOOL_RESULT_THRESHOLD
        passes_research_trim = is_str and name in _RESEARCH_TOOL_NAMES and clen > _RESEARCH_TRIM_MIN_CHARS
        if passes_research_trim:
            research_indices.append(i)
        elif passes_large:
            other_indices.append(i)

    if research_retain is not None:
        for i, msg in enumerate(messages):
            if not isinstance(msg, dict):
                continue
            content = tool_result_text_for_trim(msg)
            if not isinstance(content, str):
                continue
            name = tool_call_id_to_name.get(msg.get("tool_call_id") or "")
            if not name:
                name = xml_result_index_to_name.get(i, "")
            if name == "append_research_scratch" or "Scratch file:" in content:
                if content not in research_retain and not content.strip().startswith("[Research result from an earlier round") and not content.strip().startswith("[Large tool result"):
                    research_retain.append(content)

    _ov_chars = (
        other_placeholder_preview_chars
        if other_placeholder_preview_chars is not None
        else 120
    )
    _rf_chars = (
        research_placeholder_flat_chars if research_placeholder_flat_chars is not None else 300
    )
    _ov_chars = max(16, min(400, _ov_chars))
    _rf_chars = max(32, min(800, _rf_chars))

    _other_keep_eff = max(0, _LARGE_TOOL_RESULT_KEEP_LAST_N - max(0, int(other_keep_subtract or 0)))
    _other_take_count = len(other_indices) - _other_keep_eff
    if _other_take_count <= 0:
        _other_target_indices: list[int] = []
    else:
        _other_target_indices = other_indices[:_other_take_count]

    for i in _other_target_indices:
        if _pins and str(messages[i].get("tool_call_id") or "") in _pins:
            continue
        orig = messages[i].get("content") or ""
        if not isinstance(orig, str):
            continue
        if len(orig) < _LARGE_TOOL_RESULT_THRESHOLD:
            continue
        preview = (orig[:200] + "…") if len(orig) > 200 else orig
        _flat = preview[:_ov_chars].replace(chr(10), " ")
        new_content = (
            f"[Large tool result omitted to reduce request size ({len(orig):,} chars). "
            f"Preview: {_flat}. Use the latest tool results above to continue."
            f"{_stable_placeholder_appendix(orig)}]"
        )
        messages[i] = {**messages[i], "content": new_content}
        deltas.append((i, len(orig), len(new_content)))

    _placeholder_prefix = "[Research result from an earlier round"
    _large_placeholder_prefix = "[Large tool result"
    keep_n = research_keep_last_n if research_keep_last_n is not None else _RESEARCH_LARGE_RESULT_KEEP_LAST_N
    for i in research_indices[: -keep_n]:
        orig = messages[i].get("content") or ""
        if not isinstance(orig, str):
            continue
        tool_name = tool_call_id_to_name.get(messages[i].get("tool_call_id") or "", "")
        if tool_name == "append_research_scratch" or "Scratch file:" in orig:
            if research_retain is not None and not orig.strip().startswith(_placeholder_prefix) and not orig.strip().startswith(_large_placeholder_prefix):
                research_retain.append(orig)
            continue
        if research_retain is not None and not orig.strip().startswith(_placeholder_prefix) and not orig.strip().startswith(_large_placeholder_prefix):
            research_retain.append(orig)
        if _pins and str(messages[i].get("tool_call_id") or "") in _pins:
            continue
        preview = (orig[:400] + "…") if len(orig) > 400 else orig
        preview_flat = preview[:_rf_chars].replace(chr(10), " ")
        placeholder = _LARGE_RESEARCH_PLACEHOLDER_BASE + f"Preview: {preview_flat}{_stable_placeholder_appendix(orig)}]"
        if len(orig) <= len(placeholder):
            continue
        messages[i] = {**messages[i], "content": placeholder}
        deltas.append((i, len(orig), len(placeholder)))
    return deltas


# Re-export for run_raw_agent and tests (same names as before).
RESEARCH_INDICATOR_TOOLS = _RESEARCH_INDICATOR_TOOLS
RESEARCH_DEPTH_TOOLS = _RESEARCH_DEPTH_TOOLS
SAVE_RESEARCH_REMINDER_ROUND = _SAVE_RESEARCH_REMINDER_ROUND
CTX_SMALL_THRESHOLD_TOKENS = _CTX_SMALL_THRESHOLD_TOKENS
