#!/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

    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 = ""
    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])
            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)

    # Escalation logic
    if success:
        fstat = "done"
    elif retry_count >= 4:
        fstat = "done"
        outcome = f"FAILED (5 attempts): {outcome}"
        log.warning("Poison-pill: task %s failed after %d retries", 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)

    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)
