#!/usr/bin/env python3
# Copyright (c) 2026 PULSE Contributors. Proprietary. All rights reserved.
"""LLM tool loop execution engine."""
from __future__ import annotations

import logging
from pathlib import Path

from pulse.core.brain_utils import TASK_BOARD_FILE, now_iso as _now_iso
from pulse.tasks.task_ops import update_registry, finalize_task as _finalize_task_shared

from pulse.workers.worker_tools import tool_brain_query, compute_salience, capture_to_bridge_log, save_session_json
from pulse.workers.runner.worker_context import build_system_context, load_domain_skills, save_learnings


def execute_task(task: dict, worker_id: str, model: str, tier: str, provider: str, skills_used: bool, llm, log: logging.Logger, tool_names: list[str], max_tool_iterations: int, send_heartbeat_cb, log_dir: Path):
    """Run the main LLM tool loop for a single claimed task."""
    task_id = task.get("id") or task.get("task_id", "unknown")
    title, domain = task.get("title", "Unknown task"), task.get("domain", "general")
    description, retry_count = task.get("description", title), task.get("retry_count", 0)

    # Layer 11: Inline compliance gate
    from pulse.workers.runner.worker_compliance import check_task_compliance
    passed, reason = check_task_compliance(task)
    if not passed:
        log.warning("COMPLIANCE REJECT: [%s] %s — %s", task_id, title, reason)
        _finalize_task_shared(TASK_BOARD_FILE, task_id, worker_id,
                              outcome=f"Rejected: {reason}", status="rejected")
        update_registry(worker_id, status="idle")
        return

    # Dead-man's switch: refuse execution if compliance daemon is offline
    try:
        import json as _json
        _pids_file = Path(__file__).resolve().parent.parent.parent.parent / "state" / "daemon_pids.json"
        if _pids_file.exists():
            _pids = _json.loads(_pids_file.read_text(encoding="utf-8"))
            if "compliance_daemon" not in _pids:
                log.error("DEAD_MAN_SWITCH: Compliance daemon not registered. Refusing task.")
                _finalize_task_shared(TASK_BOARD_FILE, task_id, worker_id,
                                      outcome="Blocked: compliance daemon offline", status="blocked")
                update_registry(worker_id, status="idle")
                return
    except Exception:
        pass  # Graceful degradation — don't block workers if PID file unreadable

    log.info("EXECUTING: [%s] %s (retry %d)", task_id, title, retry_count)
    update_registry(worker_id, status="working", task_id=task_id)
    trace: list[dict] = []

    def _rec(etype, content="", tool="", args=None, error=False):
        trace.append({
            "timestamp": _now_iso(), "type": etype, "content": str(content)[:2000],
            "tool": tool, "args": args or {}, "salience": compute_salience(etype, tool, error)
        })

    _rec("task_start", f"Claimed: {title} [{domain}]")
    
    sys_ctx = build_system_context(worker_id, tier, title, domain)
    sys_ctx += "\nExecute autonomously. Tools: " + ", ".join(tool_names) + f". Max {max_tool_iterations} iterations.\n"

    sc, loaded_skills = load_domain_skills(domain, skills_used)
    if sc:
        sys_ctx += f"\n{sc}\n"

    brain_intel = ""
    # P1-1: Collect item IDs from the initial brain_query so reward signal has accurate credit data.
    # brain_search returns results dict with per-source hit lists; each hit carries entry_id/id.
    _initial_brain_item_ids: list[str] = []
    try:
        from pulse.brain.brain_search import brain_search as _brain_search_fn
        _bq_results = _brain_search_fn(f"{title} {domain}", max_per_source=5)
        for _src_hits in _bq_results.get("results", {}).values():
            if not isinstance(_src_hits, list):
                continue
            for _hit in _src_hits:
                _eid = _hit.get("entry_id") or _hit.get("id") or ""
                if _eid and _eid not in _initial_brain_item_ids:
                    _initial_brain_item_ids.append(_eid)
    except Exception:
        pass
    try:
        brain_intel = tool_brain_query(f"{title} {domain}")
        if brain_intel and brain_intel != "No brain results.":
            _rec("brain_query", brain_intel[:500])
    except Exception:
        pass

    if brain_intel and brain_intel != "No brain results.":
        sys_ctx += f"\n\n## BRAIN KNOWLEDGE (apply this, avoid past failures)\n{brain_intel[:2000]}"

    user_msg = f"Execute this task:\n\nTitle: {title}\nDescription: {description}\nDomain: {domain}\nTask ID: {task_id}"
    if brain_intel and brain_intel != "No brain results.":
        user_msg = f"## BRAIN KNOWLEDGE\n{brain_intel}\n\nApply above knowledge. Avoid past failures.\n\n{user_msg}"

    messages = [{"role": "user", "content": user_msg}]
    outcome, failures, total_in, total_out, success = "", [], 0, 0, False

    try:
        llm.init_client()
        for iteration in range(max_tool_iterations):
            send_heartbeat_cb(task_id=task_id, iteration=iteration + 1, status="working")
            text, tool_calls, usage = llm.call_llm(sys_ctx, messages)
            total_in += usage.get("in", 0)
            total_out += usage.get("out", 0)
            if text:
                _rec("reasoning", text[:1000])
                # Conscience hook: detect ethical drift in reasoning
                try:
                    from pulse.brain.integrity.conscience_hook import audit_reasoning
                    drift = audit_reasoning(worker_id, text, iteration)
                    if drift:
                        log.warning("CONSCIENCE DRIFT: %s at iteration %d", drift["drift_type"], iteration)
                        _rec("conscience_drift", f"{drift['drift_type']}: {drift['trigger']}", error=True)
                except Exception:
                    pass
            if not tool_calls:
                outcome = text or "No output."
                success = False
                break
            completed = any(tc["name"] == "complete_task" for tc in tool_calls)
            if completed:
                outcome = next(tc["input"].get("outcome", "Done.") for tc in tool_calls if tc["name"] == "complete_task")
                success = True
            for tc in tool_calls:
                _rec("tool_call", tool=tc["name"], args=tc.get("input", {}))
            results = llm.run_tools(tool_calls)
            
            # Brain-first enforcement
            for i, tc in enumerate(tool_calls):
                if tc["name"] in ("write_file", "edit_file"):
                    path = tc.get("input", {}).get("path", "")
                    try:
                        brain_hint = tool_brain_query(f"writing to {path}")
                        if brain_hint and brain_hint != "No brain results.":
                            results[i] = f"[BRAIN] {brain_hint[:400]}\n\n{results[i]}"
                            _rec("brain_query", f"[auto-write] {brain_hint[:300]}", tool="brain_query")
                    except Exception:
                        pass

            for tc, result in zip(tool_calls, results):
                is_err = any(kw in str(result).lower() for kw in ("error", "traceback", "failed", "exception", "permission denied"))
                _rec("tool_result", str(result)[:500], tool=tc["name"], error=is_err)
                if is_err:
                    _rec("error", f"{tc['name']}: {str(result)[:300]}", error=True)
            
            llm.append_tool_results(messages, text, tool_calls, results)
            if completed:
                break
        else:
            outcome = "Hit max iterations."
            failures.append(f"Max iterations ({max_tool_iterations})")
            _rec("error", outcome, error=True)
    except Exception as e:
        outcome, failures = f"Error: {e}", [str(e)]
        _rec("error", str(e), error=True)
        log.error("Task error: %s", e, exc_info=True)

    _rec("complete", outcome[:500])
    rate = 0.0 if provider == "ollama" else (0.15 if "flash" in model.lower() else 1.25)
    cost = (total_in + total_out) * (rate / 1_000_000.0)

    try:
        from pulse.daemons.neural.neuromodulation import get_patience_limit, update_patience
        patience_limit = get_patience_limit(worker_id)
        update_patience(worker_id, success)
    except Exception:
        patience_limit = 4  # Fallback to hardcoded

    # Escalation logic
    if success:
        fstat = "done"
    elif retry_count >= patience_limit:
        fstat = "done"
        outcome = f"FAILED ({retry_count + 1} attempts): {outcome}"   
        log.warning("Poison-pill: task %s failed after %d retries (5-HT exhausted)", task_id, retry_count + 1)
    elif retry_count >= 1:
        fstat = "escalated"
    else:
        fstat = "open"
    _finalize_task_shared(
        TASK_BOARD_FILE, task_id, worker_id,
        outcome=outcome, status=fstat, tokens_in=total_in, tokens_out=total_out, cost=cost
    )
    capture_to_bridge_log(worker_id, trace, task_id,
                          task_title=task.get("title", ""))
    save_session_json(worker_id, model, tier, task_id, task, trace, outcome, success, total_in, total_out, cost, skills_used=loaded_skills)

    # Phase 4: Record task outcome for reward signal pipeline
    # Populates credit_ledger → activates learned prompt + procedural cache
    # P1-1 FIX: Use item IDs collected from the initial brain_search call (before task loop)
    # instead of scraping FTS from trace text — that heuristic was ~92.5% blind.
    try:
        from pulse.core.reward.reward_signal import record_task_outcome
        record_task_outcome(
            task_id=task_id, agent=worker_id, success=success,
            domain=domain, retrieved_item_ids=_initial_brain_item_ids[:20]
        )
    except Exception as e:
        log.debug("Reward signal failed: %s", e)

    save_learnings(log, worker_id, task_id, outcome, domain, trace, failures, loaded_skills)

    update_registry(worker_id, status="idle", domain=domain)
    
    with open(log_dir / f"{worker_id}.log", "a", encoding="utf-8") as f:
        f.write(f"\n[{_now_iso()}] {task_id} | {fstat} | {len(trace)} events\nOutcome: {outcome[:200]}\n")
    log.info("DONE: [%s] %s ($%.4f) | %s", task_id, outcome[:100], cost, fstat)
