"""Base AI Agent — common functionality for all agents.

Each agent:
- Has a name, role, and specialization
- Runs in a background thread (persistent)
- Picks up goals from the goal memory
- Uses the LLM (via harness) to generate responses
- Reports progress and results
- Can be paused/resumed
- Tracks stats (goals completed, steps executed, errors)
"""

from __future__ import annotations

import logging
import threading
import time
from dataclasses import dataclass, field
from typing import Any, Callable

from ..memory.goal_memory import GoalMemory, Goal
from ..memory.persistent import PersistentMemory

logger = logging.getLogger(__name__)


@dataclass
class AgentStats:
    """Stats for a single agent."""
    goals_assigned: int = 0
    goals_completed: int = 0
    goals_failed: int = 0
    steps_executed: int = 0
    errors: int = 0
    uptime_s: float = 0.0
    last_active: float = 0.0


class BaseAgent:
    """Base class for all AI agents.

    Subclasses implement `process_goal()` which is called when a goal
    is assigned to this agent.
    """

    def __init__(self, name: str, role: str, description: str,
                 goal_memory: GoalMemory,
                 persistent_memory: PersistentMemory | None = None,
                 generate_fn: Callable[[str], str] | None = None,
                 poll_interval_s: float = 5.0) -> None:
        self.name = name
        self.role = role
        self.description = description
        self.goal_memory = goal_memory
        self.persistent_memory = persistent_memory
        self._generate_fn = generate_fn
        self._poll_interval = poll_interval_s

        self._running = False
        self._paused = False
        self._thread: threading.Thread | None = None
        self._current_goal: Goal | None = None
        self._stats = AgentStats()
        self._start_time = 0.0

    def set_generate_fn(self, fn: Callable[[str], str]) -> None:
        """Set the function used to generate LLM responses."""
        self._generate_fn = fn

    def start(self) -> None:
        """Start the agent in a background thread."""
        if self._running:
            return
        self._running = True
        self._paused = False
        self._start_time = time.time()
        self._thread = threading.Thread(target=self._run_loop, daemon=True, name=f"agent-{self.name}")
        self._thread.start()
        logger.info("Agent '%s' started (%s)", self.name, self.role)

    def stop(self) -> None:
        """Stop the agent."""
        self._running = False
        if self._thread:
            self._thread.join(timeout=10)
        logger.info("Agent '%s' stopped", self.name)

    def pause(self) -> None:
        """Pause the agent (doesn't pick up new goals)."""
        self._paused = True
        logger.info("Agent '%s' paused", self.name)

    def resume(self) -> None:
        """Resume the agent."""
        self._paused = False
        logger.info("Agent '%s' resumed", self.name)

    def _run_loop(self) -> None:
        """Main agent loop — continuously picks up and processes goals."""
        while self._running:
            try:
                if self._paused:
                    time.sleep(self._poll_interval)
                    continue

                # Find a goal assigned to this agent
                goals = self.goal_memory.get_goals_for_agent(self.name)
                active = [g for g in goals if g.status in ("in_progress", "planning")]

                if not active:
                    # Try to pick up an unassigned goal that matches our role
                    pending = self.goal_memory.list_goals(status="pending")
                    for g in pending:
                        if self._can_handle(g):
                            self.goal_memory.assign_agent(g.id, self.name)
                            self._stats.goals_assigned += 1
                            active = [g]
                            break

                if not active:
                    time.sleep(self._poll_interval)
                    continue

                # Process the first active goal
                goal = active[0]
                self._current_goal = goal
                self._stats.last_active = time.time()

                result = self.process_goal(goal)

                if result.get("success"):
                    self._stats.steps_executed += 1
                    # Advance the goal
                    self.goal_memory.execute_step(
                        goal.id, result.get("output", ""), success=True
                    )
                else:
                    self._stats.errors += 1
                    self.goal_memory.execute_step(
                        goal.id, result.get("error", "Unknown error"), success=False
                    )

                # Check if goal is completed
                updated = self.goal_memory.get_goal(goal.id)
                if updated and updated.status == "completed":
                    self._stats.goals_completed += 1
                    if self.persistent_memory:
                        self.persistent_memory.add_episodic(
                            "event", f"Agent {self.name} completed goal: {goal.title}",
                            importance=0.8, tags=["goal", "completed", self.name]
                        )
                    logger.info("Agent '%s' completed goal: %s", self.name, goal.title)

                self._current_goal = None

            except Exception as e:
                logger.error("Agent '%s' error: %s", self.name, e)
                self._stats.errors += 1
                time.sleep(self._poll_interval)

            self._stats.uptime_s = time.time() - self._start_time

    def _can_handle(self, goal: Goal) -> bool:
        """Check if this agent can handle a goal. Override in subclasses."""
        return True

    def process_goal(self, goal: Goal) -> dict[str, Any]:
        """Process a goal step. Must be implemented by subclasses.

        Returns: {"success": bool, "output": str, "error": str}
        """
        raise NotImplementedError

    def _generate(self, prompt: str) -> str:
        """Generate a response using the LLM."""
        if self._generate_fn:
            return self._generate_fn(prompt)
        return "(no generation function available)"

    def get_status(self) -> dict[str, Any]:
        """Get agent status."""
        return {
            "name": self.name,
            "role": self.role,
            "description": self.description,
            "running": self._running,
            "paused": self._paused,
            "current_goal": self._current_goal.title if self._current_goal else None,
            "stats": {
                "goals_assigned": self._stats.goals_assigned,
                "goals_completed": self._stats.goals_completed,
                "goals_failed": self._stats.goals_failed,
                "steps_executed": self._stats.steps_executed,
                "errors": self._stats.errors,
                "uptime_s": round(self._stats.uptime_s, 1),
                "last_active": self._stats.last_active,
            },
        }