"""
GenWorker Desktop Agent — v11 (Modular)

Uses Claude Computer Use API with three Anthropic tools:
  - computer: mouse / keyboard / screenshot
  - bash: shell command execution (PowerShell on Windows, bash/zsh on macOS/Linux)
  - str_replace_based_edit_tool: file view / create / edit

Usage:
    genworker                                     # interactive prompt
    genworker "Open Notepad and write Hello World"
    genworker --list-monitors
    genworker --monitor 1 "Open Chrome"
    genworker --no-thinking "Simple task"
    genworker --timeout 300 "Quick task"
    genworker --no-grid "Task without coordinate grid"
    genworker --gui                               # launch optional GUI
"""

import argparse
import sys
import time

import anthropic

from genworker.config import (
    MODEL, MAX_STEPS, MAX_CONTEXT_IMAGES, TASK_TIMEOUT, MAX_NUDGES,
    SCREENSHOT_GRID_OVERLAY, DEFAULT_MONITOR,
    client, log,
)
import genworker.config as cfg
import genworker.display as ui
from genworker.context import compact_context, estimate_context_size, ErrorTracker
from genworker.prompts import SYSTEM_PROMPT, preprocess_query, get_tools
from genworker.tools import execute_tool
from genworker.tools.computer import (
    capture_screen, get_screenshot_media_type,
    list_monitors, get_monitor, diagnose_macos_mouse,
)


# ── Tool emoji map ─────────────────────────────────────────────────────────────

_TOOL_EMOJI = {
    "computer":                   "🖱 ",
    "bash":                       "💻",
    "str_replace_based_edit_tool":"📝",
    "str_replace_editor":         "📝",
}


# ═══════════════════════════════════════════════════════════════════════════════
# MAIN AGENT LOOP
# ═══════════════════════════════════════════════════════════════════════════════

def run_agent(
    task: str,
    monitor_index: int = DEFAULT_MONITOR,
    enable_thinking: bool = True,
    timeout: int = TASK_TIMEOUT,
):
    monitor  = get_monitor(monitor_index)
    grid_on  = cfg.SCREENSHOT_GRID_OVERLAY

    # macOS: run mouse diagnostics and auto-fix Retina scaling if needed
    if cfg.IS_MACOS:
        monitor = diagnose_macos_mouse(monitor)

    log.info(f"Monitor [{monitor_index}]: {monitor['name']} {monitor['width']}x{monitor['height']}")
    log.info(f"Task: {task}")

    ui.banner(
        version      = "v11",
        model        = MODEL,
        monitor_index= monitor_index,
        monitor_name = monitor["name"],
        monitor_res  = f"{monitor['width']}×{monitor['height']}",
        thinking     = enable_thinking,
        grid         = grid_on,
        timeout      = timeout,
        max_steps    = MAX_STEPS,
        task         = task,
    )

    start_time = time.time()

    # ── Step 0: Preprocess query ───────────────────────────────────────────
    execution_plan = preprocess_query(task)

    ui.countdown(3)

    # ── Initial message ────────────────────────────────────────────────────
    media_type = get_screenshot_media_type()
    image_b64  = capture_screen(monitor)
    messages   = [
        {
            "role": "user",
            "content": [
                {
                    "type": "text",
                    "text": (
                        f"## TASK\n{task}\n\n"
                        f"## EXECUTION PLAN\n{execution_plan}\n\n"
                        f"Here is the current screen. Begin executing the plan. "
                        f"Review the objectives, then work through them one by one. "
                        f"Remember: prefer bash/str_replace_based_edit_tool over GUI whenever possible."
                    ),
                },
                {
                    "type": "image",
                    "source": {"type": "base64", "media_type": media_type, "data": image_b64},
                },
            ],
        }
    ]

    tools                 = get_tools()
    step                  = 0
    nudge_count           = 0
    error_tracker         = ErrorTracker()
    api_retries           = 0
    total_input_tokens    = 0
    total_output_tokens   = 0

    while step < MAX_STEPS:
        step += 1

        # ── Check timeout ──────────────────────────────────────────────────
        elapsed = time.time() - start_time
        if elapsed > timeout:
            ui.print_timeout(timeout, elapsed)
            log.warning(f"Task timeout after {elapsed:.0f}s, {step} steps")
            messages.append({
                "role": "user",
                "content": [{
                    "type": "text",
                    "text": (
                        f"TIMEOUT: You have {timeout}s and time is up. "
                        f"Summarize what objectives were completed and which remain. "
                        f"Say TASK_COMPLETE if all done, otherwise list what's left."
                    ),
                }],
            })
            try:
                response = client.beta.messages.create(
                    model=MODEL, max_tokens=2048, system=SYSTEM_PROMPT,
                    tools=tools, messages=messages, betas=["computer-use-2025-01-24"],
                )
                for block in response.content:
                    if hasattr(block, "text"):
                        ui.print_claude(block.text)
            except Exception:
                pass
            break

        ui.step_header(step, MAX_STEPS, elapsed)

        # ── Adaptive context compaction ────────────────────────────────────
        ctx = estimate_context_size(messages)
        if ctx["est_tokens"] > 80000 or ctx["image_count"] > MAX_CONTEXT_IMAGES:
            log.info(f"Context: ~{ctx['est_tokens']} tokens, {ctx['image_count']} images — compacting")
            ui.print_context_compact(ctx["image_count"])
            messages = compact_context(messages, max_images=max(5, MAX_CONTEXT_IMAGES - 5))
        elif step % 15 == 0:
            messages = compact_context(messages)

        # ── Stuck-loop detection ───────────────────────────────────────────
        if error_tracker.is_stuck_in_loop():
            log.warning("Agent appears stuck in a loop")
            ui.print_stuck_loop()
            messages.append({
                "role": "user",
                "content": [{
                    "type": "text",
                    "text": (
                        "⚠️ You appear to be stuck repeating the same actions. STOP and:\n"
                        "1. Take a screenshot to see the current state\n"
                        "2. Re-read the execution plan\n"
                        "3. Try a COMPLETELY DIFFERENT approach\n"
                        "4. If GUI is failing, switch to bash/CLI\n"
                    ),
                }],
            })
            error_tracker.consecutive_errors = 0
            error_tracker.last_actions = []

        try:
            # ── Call Claude ────────────────────────────────────────────────
            response = client.beta.messages.create(
                model=MODEL,
                max_tokens=8192,
                system=SYSTEM_PROMPT,
                tools=tools,
                messages=messages,
                betas=["computer-use-2025-01-24"],
            )
            api_retries = 0

            if hasattr(response, "usage"):
                total_input_tokens  += getattr(response.usage, "input_tokens", 0)
                total_output_tokens += getattr(response.usage, "output_tokens", 0)

            log.info(
                f"Stop reason: {response.stop_reason} | "
                f"Tokens in/out: {total_input_tokens}/{total_output_tokens}"
            )

            messages.append({"role": "assistant", "content": response.content})

            # ── Process response blocks ────────────────────────────────────
            tool_results = []
            task_done    = False

            for block in response.content:

                if hasattr(block, "type") and block.type == "thinking":
                    thinking_text = getattr(block, "thinking", "")
                    ui.print_thinking(thinking_text)
                    log.info(f"Thinking: {thinking_text[:500]}")

                elif block.type == "text":
                    ui.print_claude(block.text)
                    log.info(f"Claude: {block.text}")
                    if "TASK_COMPLETE" in block.text:
                        task_done = True

                elif block.type == "tool_use":
                    result_text = execute_tool(block, monitor)

                    if result_text.startswith("Error") or result_text.startswith("⛔"):
                        error_tracker.record_error(
                            f"{block.name}/{block.input.get('action', block.input.get('command', '')[:50])}",
                            result_text[:100],
                        )
                    else:
                        error_tracker.record_success(
                            f"{block.name}/{block.input.get('action', block.input.get('command', '')[:30])}"
                        )

                    emoji = _TOOL_EMOJI.get(block.name, "⚡")
                    ui.print_tool(emoji, block.name, result_text)

                    tool_content = []

                    if block.name == "computer":
                        if result_text == "__screenshot__":
                            image_b64 = capture_screen(monitor)
                            tool_content.append({
                                "type": "image",
                                "source": {"type": "base64", "media_type": media_type, "data": image_b64},
                            })
                        else:
                            from genworker.tools.computer import get_adaptive_delay
                            delay     = get_adaptive_delay(block.input.get("action", ""), block.input)
                            time.sleep(delay)
                            image_b64 = capture_screen(monitor)
                            tool_content.append({"type": "text", "text": result_text})
                            tool_content.append({
                                "type": "image",
                                "source": {"type": "base64", "media_type": media_type, "data": image_b64},
                            })
                    else:
                        tool_content.append({"type": "text", "text": result_text})
                        if error_tracker.needs_recovery():
                            tool_content.append({
                                "type": "text",
                                "text": error_tracker.get_recovery_hint(),
                            })

                    tool_results.append({
                        "type":        "tool_result",
                        "tool_use_id": block.id,
                        "content":     tool_content,
                    })

            # ── Completion check ───────────────────────────────────────────
            if task_done:
                elapsed = time.time() - start_time
                ui.print_success(step, elapsed, total_input_tokens, total_output_tokens)
                log.info(
                    f"TASK_COMPLETE after {step} steps, {elapsed:.0f}s, "
                    f"{total_input_tokens}+{total_output_tokens} tokens"
                )
                break

            # ── Route on stop reason ───────────────────────────────────────
            if response.stop_reason == "end_turn":
                if tool_results:
                    messages.append({"role": "user", "content": tool_results})
                elif not task_done:
                    nudge_count += 1
                    log.warning(f"Claude stopped without TASK_COMPLETE — nudge #{nudge_count}")

                    if nudge_count >= MAX_NUDGES:
                        ui.print_max_nudges()
                        messages.append({
                            "role": "user",
                            "content": [{
                                "type": "text",
                                "text": (
                                    "IMPORTANT: You have been nudged multiple times. You MUST now either:\n"
                                    "1. Continue working on remaining objectives (take a screenshot and proceed)\n"
                                    "2. If ALL objectives are truly complete, say TASK_COMPLETE with a summary\n"
                                    "3. If you're stuck, explain what's blocking you\n"
                                    "Do NOT just acknowledge — take ACTION."
                                ),
                            }],
                        })
                    else:
                        ui.print_nudge(nudge_count)
                        image_b64 = capture_screen(monitor)
                        messages.append({
                            "role": "user",
                            "content": [
                                {
                                    "type": "text",
                                    "text": (
                                        "You stopped but have not said TASK_COMPLETE. "
                                        "Review the execution plan — are ALL objectives completed and verified? "
                                        "If any remain, continue working. Take a screenshot if needed."
                                    ),
                                },
                                {
                                    "type": "image",
                                    "source": {"type": "base64", "media_type": media_type, "data": image_b64},
                                },
                            ],
                        })

            elif response.stop_reason == "tool_use":
                if tool_results:
                    messages.append({"role": "user", "content": tool_results})
                else:
                    log.warning("tool_use with no results — sending fresh screenshot")
                    image_b64 = capture_screen(monitor)
                    messages.append({
                        "role": "user",
                        "content": [
                            {"type": "text", "text": "Continue with the task:"},
                            {"type": "image", "source": {"type": "base64", "media_type": media_type, "data": image_b64}},
                        ],
                    })

        except KeyboardInterrupt:
            ui.print_interrupted()
            break

        except anthropic.APIError as e:
            api_retries += 1
            log.error(f"API error (attempt {api_retries}): {e}")

            if api_retries >= 5:
                ui.print_api_give_up(api_retries)
                break

            wait_time = min(2 ** api_retries, 30)
            ui.print_api_error(str(e), api_retries, wait_time)
            time.sleep(wait_time)

            if "context" in str(e).lower() or "token" in str(e).lower():
                log.info("Possible context overflow — aggressive compaction")
                messages = compact_context(messages, max_images=5)
            continue

        except Exception as e:
            log.error(f"Loop error: {e}", exc_info=True)
            ui.print_tool("❌", "error", str(e))
            break

    else:
        elapsed = time.time() - start_time
        ui.print_max_steps(MAX_STEPS, elapsed)

    # ── Final stats ────────────────────────────────────────────────────────
    elapsed = time.time() - start_time
    ui.print_final_stats(step, elapsed, total_input_tokens, total_output_tokens,
                         len(error_tracker.error_history))


# ═══════════════════════════════════════════════════════════════════════════════
# CLI ENTRY POINT
# ═══════════════════════════════════════════════════════════════════════════════

def main():
    parser = argparse.ArgumentParser(
        description="GenWorker v11 — Desktop AI Agent (Claude Computer Use)"
    )
    parser.add_argument("task",          nargs="?", default=None,          help="Task to execute")
    parser.add_argument("--monitor",     "-m", type=int, default=DEFAULT_MONITOR, help="Monitor index")
    parser.add_argument("--list-monitors", "-l", action="store_true",      help="List available monitors")
    parser.add_argument("--no-thinking", action="store_true",              help="Disable extended thinking")
    parser.add_argument("--timeout",     "-t", type=int, default=TASK_TIMEOUT, help="Task timeout in seconds")
    parser.add_argument("--no-grid",     action="store_true",              help="Disable coordinate grid overlay on screenshots")
    parser.add_argument("--gui",         action="store_true",              help="Launch the optional Tkinter GUI")
    args = parser.parse_args()

    if args.list_monitors:
        list_monitors()
        sys.exit(0)

    if args.gui:
        try:
            from genworker.gui import launch_gui
            launch_gui()
        except ImportError as e:
            print(f"❌  GUI not available: {e}")
            print("    Make sure tkinter is installed (it ships with Python on most platforms).")
            sys.exit(1)
        return

    if args.no_grid:
        cfg.SCREENSHOT_GRID_OVERLAY = False

    list_monitors()

    task = args.task
    if not task:
        task = input("📝  Enter task: ").strip()
    if not task:
        task = "Open Notepad and type Hello World"

    run_agent(
        task,
        monitor_index=args.monitor,
        enable_thinking=not args.no_thinking,
        timeout=args.timeout,
    )


if __name__ == "__main__":
    main()
