"""
Computer Use Agent - Main agent loop.

Screenshot → Vision Model → Action → Execute → Repeat

Every action gets a NatraLocke-compatible receipt.
"""

import json
import time
import uuid
from datetime import datetime
from pathlib import Path
from typing import Optional, Callable

from .models import Action, ActionType, AgentStep, ComputerUseReceipt
from .screen import ScreenController
from .vision import OllamaVisionClient


class ComputerUseAgent:
    """
    Autonomous computer use agent with vision.

    Usage:
        agent = ComputerUseAgent(task="Open Firefox and go to google.com")
        receipt = agent.run(max_steps=20)
        print(receipt.success, receipt.final_reason)
    """

    def __init__(
        self,
        task: str,
        model: str = "llava:7b",
        screenshot_dir: Optional[str] = None,
        receipt_dir: Optional[str] = None,
        action_delay: float = 0.5,
        on_step: Optional[Callable[[AgentStep], None]] = None,
    ):
        self.task = task
        self.model = model
        self.on_step = on_step

        # Set up directories
        self.screenshot_dir = Path(screenshot_dir) if screenshot_dir else Path("/tmp/computer_use/screenshots")
        self.receipt_dir = Path(receipt_dir) if receipt_dir else Path("receipts/computer_use")
        self.screenshot_dir.mkdir(parents=True, exist_ok=True)
        self.receipt_dir.mkdir(parents=True, exist_ok=True)

        # Initialize components
        self.screen = ScreenController(
            action_delay=action_delay,
            screenshot_dir=str(self.screenshot_dir),
        )
        self.vision = OllamaVisionClient(model=model)

        # Session state
        self.session_id = str(uuid.uuid4())[:8]
        self.steps: list[AgentStep] = []
        self.running = False

    def run(self, max_steps: int = 30, verbose: bool = True) -> ComputerUseReceipt:
        """
        Run the agent loop until task is complete or max_steps reached.

        Args:
            max_steps: Maximum number of steps before giving up
            verbose: Print progress to stdout

        Returns:
            ComputerUseReceipt with full audit trail
        """
        # Initialize receipt
        receipt = ComputerUseReceipt(
            session_id=self.session_id,
            task_description=self.task,
            started_at=datetime.utcnow().isoformat(),
            model_name=self.model,
            screen_resolution=self.screen.screen_size,
        )

        self.running = True

        if verbose:
            print(f"[Agent] Starting task: {self.task}")
            print(f"[Agent] Model: {self.model}")
            print(f"[Agent] Screen: {receipt.screen_resolution}")
            print(f"[Agent] Max steps: {max_steps}")
            print()

        # Check vision model health
        if not self.vision.check_health():
            if verbose:
                print(f"[Agent] ERROR: Vision model '{self.model}' not available")
                print(f"[Agent] Try: ollama pull {self.model}")
            receipt.finalize(success=False, reason=f"Vision model {self.model} not available")
            return receipt

        try:
            for step_num in range(1, max_steps + 1):
                if not self.running:
                    break

                step = self._execute_step(step_num, verbose)
                self.steps.append(step)
                receipt.steps.append(step)

                if self.on_step:
                    self.on_step(step)

                # Check for termination
                if step.action:
                    if step.action.action_type == ActionType.DONE:
                        receipt.finalize(success=True, reason=step.action.reason)
                        break
                    elif step.action.action_type == ActionType.FAILED:
                        receipt.finalize(success=False, reason=step.action.reason)
                        break

            else:
                # Max steps reached
                receipt.finalize(success=False, reason=f"Max steps ({max_steps}) reached")

        except KeyboardInterrupt:
            receipt.finalize(success=False, reason="Interrupted by user")
            if verbose:
                print("\n[Agent] Interrupted")

        except Exception as e:
            receipt.finalize(success=False, reason=f"Error: {str(e)}")
            if verbose:
                print(f"\n[Agent] Error: {e}")

        # Save receipt
        receipt_path = self.receipt_dir / f"receipt_{self.session_id}.json"
        with open(receipt_path, "w") as f:
            json.dump(receipt.model_dump(), f, indent=2, default=str)

        if verbose:
            print()
            print(f"[Agent] {'SUCCESS' if receipt.success else 'FAILED'}: {receipt.final_reason}")
            print(f"[Agent] Steps: {len(receipt.steps)}, Actions: {receipt.total_actions}")
            print(f"[Agent] Receipt: {receipt_path}")

        return receipt

    def _execute_step(self, step_num: int, verbose: bool) -> AgentStep:
        """Execute a single agent step."""
        step = AgentStep(step_number=step_num)

        # 1. Take screenshot
        img_b64, img_hash, img_path = self.screen.screenshot_base64()[0], *self.screen.screenshot()[1:]
        step.screenshot_path = img_path
        step.screenshot_hash = img_hash

        if verbose:
            print(f"[Step {step_num}] Screenshot captured")

        # 2. Build context from recent steps
        context = self._build_context()

        # 3. Send to vision model
        step.prompt_sent = self.task
        start_vision = time.perf_counter()

        action, raw_response, vision_time = self.vision.analyze_screenshot(
            image_base64=img_b64,
            task=self.task,
            context=context,
        )

        step.model_response = raw_response
        step.vision_time_ms = vision_time
        step.action = action

        if verbose:
            print(f"[Step {step_num}] Vision: {vision_time:.0f}ms -> {action.to_command()}")
            if action.reason:
                print(f"[Step {step_num}] Reason: {action.reason}")

        # 4. Execute action
        if action.action_type not in (ActionType.DONE, ActionType.FAILED, ActionType.SCREENSHOT):
            start_exec = time.perf_counter()
            self._execute_action(action)
            step.execution_time_ms = (time.perf_counter() - start_exec) * 1000
            step.action_executed = True

            if verbose:
                print(f"[Step {step_num}] Executed in {step.execution_time_ms:.0f}ms")

        return step

    def _execute_action(self, action: Action) -> None:
        """Execute a single action."""
        if action.action_type == ActionType.CLICK:
            self.screen.click(action.x, action.y)
        elif action.action_type == ActionType.DOUBLE_CLICK:
            self.screen.double_click(action.x, action.y)
        elif action.action_type == ActionType.RIGHT_CLICK:
            self.screen.right_click(action.x, action.y)
        elif action.action_type == ActionType.TYPE:
            self.screen.type_text(action.text or "")
        elif action.action_type == ActionType.PRESS:
            self.screen.press(action.key or "")
        elif action.action_type == ActionType.HOTKEY:
            self.screen.hotkey(*(action.keys or []))
        elif action.action_type == ActionType.SCROLL:
            self.screen.scroll(action.scroll_amount or 0)
        elif action.action_type == ActionType.MOVE:
            self.screen.move(action.x, action.y)
        elif action.action_type == ActionType.WAIT:
            self.screen.wait(action.duration or 1.0)

    def _build_context(self, max_steps: int = 5) -> str:
        """Build context from recent steps for the vision model."""
        if not self.steps:
            return ""

        recent = self.steps[-max_steps:]
        lines = []
        for step in recent:
            if step.action:
                lines.append(f"Step {step.step_number}: {step.action.to_command()}")
                if step.action.reason:
                    lines.append(f"  Reason: {step.action.reason}")

        return "\n".join(lines)

    def stop(self) -> None:
        """Stop the agent loop."""
        self.running = False
