# Familiar - Self-Hosted AI Companion Platform
# Copyright (c) 2026 George Scott Foley
# ORCID: 0009-0006-4957-0540
# Email: Georgescottfoley@proton.me
# Licensed under the MIT License - see LICENSE file for details

"""
Task Planner — Pre-Execution Decomposition for Complex Requests

Analyzes the user's request before the tool loop begins and produces
a structured execution plan. Simple queries bypass planning entirely.
Complex multi-step tasks get a roadmap that the LLM uses as context
during tool execution.

The plan is advisory, not rigid — the LLM still makes decisions each
iteration, but with a clear roadmap instead of stumbling through tools.

Usage:
    planner = TaskPlanner(provider)

    if planner.should_plan(message, tool_schemas):
        plan = planner.create_plan(message, tool_schemas, memory_context)
        system_prompt += plan.to_prompt()

    # Tool loop proceeds normally, but LLM has the plan as context

Integration:
    Inserted between system prompt construction and the tool loop
    in Agent.chat(). Costs one additional LLM call for complex requests.
"""

import json
import logging
import re
import time
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional

logger = logging.getLogger(__name__)


# ════════════════════════════════════════════════════════════════
# Plan Data Structures
# ════════════════════════════════════════════════════════════════


@dataclass
class PlanStep:
    """A single step in a task plan."""

    id: int
    description: str
    tool: Optional[str] = None  # Tool name, or None for LLM-only reasoning
    depends_on: List[int] = field(default_factory=list)
    status: str = "pending"  # pending, active, done, skipped

    def to_dict(self) -> dict:
        d = {"id": self.id, "description": self.description}
        if self.tool:
            d["tool"] = self.tool
        if self.depends_on:
            d["depends_on"] = self.depends_on
        return d


@dataclass
class TaskPlan:
    """A structured execution plan for a complex task."""

    goal: str
    steps: List[PlanStep]
    estimated_iterations: int = 5
    created_at: float = field(default_factory=time.time)
    planning_time_ms: float = 0.0

    def to_prompt(self) -> str:
        """Format the plan for injection into the system prompt."""
        lines = [
            "<execution_plan>",
            f"Goal: {self.goal}",
            "",
            "Steps:",
        ]
        for step in self.steps:
            tool_hint = f" [use: {step.tool}]" if step.tool else ""
            deps = (
                f" (after step {', '.join(str(d) for d in step.depends_on)})"
                if step.depends_on
                else ""
            )
            lines.append(f"  {step.id}. {step.description}{tool_hint}{deps}")

        lines.append("")
        lines.append(
            "Follow this plan as a guide. Adjust if tool results suggest a better approach. "
            "Execute steps in order, respecting dependencies."
        )
        lines.append("</execution_plan>")
        return "\n".join(lines)

    def to_dict(self) -> dict:
        return {
            "goal": self.goal,
            "steps": [s.to_dict() for s in self.steps],
            "estimated_iterations": self.estimated_iterations,
        }

    @classmethod
    def from_dict(cls, data: dict) -> "TaskPlan":
        steps = [
            PlanStep(
                id=s["id"],
                description=s["description"],
                tool=s.get("tool"),
                depends_on=s.get("depends_on", []),
            )
            for s in data.get("steps", [])
        ]
        return cls(
            goal=data.get("goal", ""),
            steps=steps,
            estimated_iterations=data.get("estimated_iterations", len(steps) + 1),
        )


# ════════════════════════════════════════════════════════════════
# Complexity Detection (lightweight, no LLM call)
# ════════════════════════════════════════════════════════════════

# Patterns that suggest multi-step tasks
_MULTI_STEP_PATTERNS = [
    re.compile(p, re.IGNORECASE)
    for p in [
        r"\b(prepare|create|build|generate|compile)\b.*(report|packet|summary|document|overview)",
        r"\b(and then|after that|once .* done|first .* then)",
        r"\b(compare|analyze|cross.?reference)\b.*\b(with|against|between)",
        r"\b(for each|for every|all of the|across all)\b",
        r"\b(step.?by.?step|multi.?step|workflow)\b",
        r"\b(import|export|migrate|transfer)\b.*\b(from|to|into)\b",
        r"\b(set up|configure|initialize|deploy)\b.*\b(with|including|and)\b",
        r"\b(review|audit|check)\b.*\b(all|every|each)\b",
        # List-like requests (commas with multiple items)
        r",\s*(and\s+)?(?:also|plus|as well as)",
        r",\s*\w+.{3,},\s*(and\s+)?\w+",  # "X, Y, and Z" pattern
    ]
]

# Patterns that suggest simple tasks (override multi-step)
_SIMPLE_PATTERNS = [
    re.compile(p, re.IGNORECASE)
    for p in [
        r"^(what|who|where|when|how much|how many)\b",
        r"^(list|show|get|find|search|check|tell me)\b",
        r"^(hi|hello|hey|thanks|thank you)",
        r"^(yes|no|ok|sure|cancel|stop|quit)",
    ]
]


def estimate_complexity(message: str, tool_count: int = 0) -> float:
    """
    Fast complexity estimate without an LLM call.

    Returns 0.0-1.0 where higher = more complex.
    Used to decide whether planning is worth the cost.
    """
    score = 0.3  # Base

    # Length factor
    words = len(message.split())
    if words > 50:
        score += 0.15
    elif words > 20:
        score += 0.1
    elif words > 10:
        score += 0.05

    # Multi-step patterns
    multi_hits = sum(1 for p in _MULTI_STEP_PATTERNS if p.search(message))
    score += min(0.4, multi_hits * 0.15)

    # Simple patterns pull score down
    if any(p.search(message) for p in _SIMPLE_PATTERNS):
        score -= 0.25

    # Many available tools = more room for multi-step
    if tool_count > 10:
        score += 0.05

    # Sentence count (multiple sentences often = multiple intentions)
    sentences = len([s for s in re.split(r"[.!?]+", message) if s.strip()])
    if sentences > 2:
        score += 0.1

    # Comma-separated items in a single request (strong multi-step signal)
    comma_items = len(re.findall(r",", message))
    if comma_items >= 2:
        score += 0.1

    return max(0.0, min(1.0, score))


# ════════════════════════════════════════════════════════════════
# Planning Prompt
# ════════════════════════════════════════════════════════════════

PLANNING_SYSTEM = """You are a task planner for an AI assistant with access to tools.
Given a user request and available tools, produce a concise execution plan.

Respond with ONLY a JSON object (no markdown, no explanation):
{
  "goal": "one-sentence summary of what the user wants",
  "steps": [
    {
      "id": 1,
      "description": "what this step accomplishes",
      "tool": "tool_name or null if reasoning-only",
      "depends_on": []
    }
  ],
  "estimated_iterations": 5
}

Rules:
- Keep plans concise: 2-8 steps maximum
- Only reference tools that exist in the available tools list
- Use null for tool when the step is thinking/summarizing (no tool needed)
- Set depends_on to step IDs that must complete first
- estimated_iterations = number of LLM calls needed (steps + final response)
- Do NOT create a step for "respond to user" or "summarize results" — that happens automatically
- Do NOT split a single tool call into multiple steps unless the second depends on the first's output
- If the request is ambiguous, plan for the most likely interpretation and note the assumption in the goal
- Group related lookups into as few steps as possible
"""


def _build_planning_message(
    user_message: str,
    tool_names: List[str],
    memory_context: str = "",
) -> str:
    """Build the message sent to the LLM for planning."""
    parts = [f"User request: {user_message}"]

    if memory_context:
        parts.append(f"\nRelevant context:\n{memory_context}")

    # Group tools by prefix when there are many to reduce token waste
    if len(tool_names) > 15:
        groups: dict = {}
        ungrouped = []
        for name in tool_names:
            prefix = name.split("_")[0] if "_" in name else ""
            if prefix:
                groups.setdefault(prefix, []).append(name)
            else:
                ungrouped.append(name)

        tool_lines = []
        for prefix, names in sorted(groups.items()):
            if len(names) >= 3:
                tool_lines.append(f"  {prefix}_*: {', '.join(names)}")
            else:
                ungrouped.extend(names)
        if ungrouped:
            tool_lines.append(f"  other: {', '.join(ungrouped)}")

        parts.append("\nAvailable tools:\n" + "\n".join(tool_lines))
    else:
        parts.append(f"\nAvailable tools: {', '.join(tool_names)}")

    return "\n".join(parts)


# ════════════════════════════════════════════════════════════════
# Task Planner
# ════════════════════════════════════════════════════════════════


class TaskPlanner:
    """
    Pre-execution planning for complex tasks.

    Analyzes the user's request and produces a structured plan
    before the agent begins calling tools. Simple queries bypass
    planning entirely.
    """

    # Complexity threshold — only plan above this
    PLANNING_THRESHOLD = 0.55

    def __init__(
        self,
        provider: Any = None,  # LLMProvider
        threshold: float = PLANNING_THRESHOLD,
        max_planning_tokens: int = 1024,
        enabled: bool = True,
    ):
        """
        Args:
            provider: LLM provider for generating plans. If None, planning is disabled.
            threshold: Complexity threshold (0-1). Only plan above this.
            max_planning_tokens: Max tokens for planning LLM call.
            enabled: Master switch to disable planning entirely.
        """
        self.provider = provider
        self.threshold = threshold
        self.max_planning_tokens = max_planning_tokens
        self.enabled = enabled

        # Stats
        self.plans_created = 0
        self.plans_skipped = 0
        self.total_planning_time_ms = 0.0

    def should_plan(
        self,
        message: str,
        tool_schemas: Optional[List[dict]] = None,
    ) -> bool:
        """
        Decide whether to create a plan for this request.

        Planning is skipped for:
        - Simple queries (greetings, single-fact lookups)
        - When no provider is available
        - When planning is disabled
        - When there are very few tools (< 3)

        Returns:
            True if planning would improve the response
        """
        if not self.enabled or not self.provider:
            return False

        tool_count = len(tool_schemas) if tool_schemas else 0
        if tool_count < 3:
            return False

        complexity = estimate_complexity(message, tool_count)
        should = complexity >= self.threshold

        if not should:
            self.plans_skipped += 1

        logger.debug(
            f"Planning decision: complexity={complexity:.2f} "
            f"threshold={self.threshold} plan={'yes' if should else 'no'}"
        )
        return should

    def create_plan(
        self,
        message: str,
        tool_schemas: List[dict],
        memory_context: str = "",
    ) -> Optional[TaskPlan]:
        """
        Generate an execution plan for the user's request.

        Makes one LLM call to decompose the task into steps.
        Returns None if planning fails (agent falls back to direct execution).

        Args:
            message: User's original message
            tool_schemas: List of tool schema dicts (with name, description)
            memory_context: Relevant memory context string

        Returns:
            TaskPlan or None on failure
        """
        if not self.provider:
            return None

        start = time.time()

        # Extract tool names for the planning prompt
        tool_names = [t.get("name", "") for t in tool_schemas if t.get("name")]

        planning_msg = _build_planning_message(message, tool_names, memory_context)

        try:
            response = self.provider.chat(
                messages=[{"role": "user", "content": planning_msg}],
                tools=[],  # No tools during planning
                system=PLANNING_SYSTEM,
                max_tokens=self.max_planning_tokens,
            )

            elapsed_ms = (time.time() - start) * 1000

            # Parse the JSON plan from response
            plan = self._parse_plan(response.text, tool_names)
            if plan:
                plan.planning_time_ms = elapsed_ms
                self.plans_created += 1
                self.total_planning_time_ms += elapsed_ms
                logger.info(
                    f"Plan created: {len(plan.steps)} steps, "
                    f"{elapsed_ms:.0f}ms, goal='{plan.goal[:60]}'"
                )
                return plan

            logger.warning("Failed to parse planning response")
            return None

        except Exception as e:
            elapsed_ms = (time.time() - start) * 1000
            logger.warning(f"Planning failed ({elapsed_ms:.0f}ms): {e}")
            return None

    def _parse_plan(
        self,
        response_text: str,
        valid_tools: List[str],
    ) -> Optional[TaskPlan]:
        """
        Parse the LLM's JSON response into a TaskPlan.

        Handles common LLM response issues:
        - Markdown code fences around JSON
        - Extra text before/after JSON
        - Tool names that don't exist in the schema
        """
        if not response_text:
            return None

        # Strip markdown code fences
        text = response_text.strip()
        text = re.sub(r"^```(?:json)?\s*", "", text)
        text = re.sub(r"\s*```$", "", text)

        # Try to extract JSON object
        try:
            data = json.loads(text)
        except json.JSONDecodeError:
            # Try to find JSON within the response
            match = re.search(r"\{[\s\S]*\}", text)
            if match:
                try:
                    data = json.loads(match.group())
                except json.JSONDecodeError:
                    return None
            else:
                return None

        if not isinstance(data, dict) or "steps" not in data:
            return None

        # Validate and clean tool references
        valid_set = set(valid_tools)
        steps = []
        for s in data["steps"]:
            if not isinstance(s, dict):
                continue
            tool = s.get("tool")
            if tool and tool not in valid_set:
                tool = None  # Remove invalid tool reference

            steps.append(
                PlanStep(
                    id=s.get("id", len(steps) + 1),
                    description=s.get("description", ""),
                    tool=tool,
                    depends_on=s.get("depends_on", []),
                )
            )

        if not steps:
            return None

        # Cap at 8 steps
        steps = steps[:8]

        return TaskPlan(
            goal=data.get("goal", ""),
            steps=steps,
            estimated_iterations=data.get("estimated_iterations", len(steps) + 1),
        )

    def get_stats(self) -> Dict[str, Any]:
        """Return planning statistics."""
        total = self.plans_created + self.plans_skipped
        return {
            "plans_created": self.plans_created,
            "plans_skipped": self.plans_skipped,
            "total_requests": total,
            "plan_rate": self.plans_created / total if total > 0 else 0,
            "avg_planning_time_ms": (
                self.total_planning_time_ms / self.plans_created if self.plans_created > 0 else 0
            ),
        }


# ════════════════════════════════════════════════════════════════
# Plan Tracker — Runtime Progress Tracking
# ════════════════════════════════════════════════════════════════


class PlanTracker:
    """
    Tracks which steps in a TaskPlan have been completed during
    the agent tool loop.

    Created when a plan is generated, updated after each tool call,
    and queried to inject progress context into the LLM prompt.
    """

    def __init__(self, plan: TaskPlan):
        self.plan = plan
        self._done: Dict[int, bool] = {}

    def record_tool_call(self, tool_name: str, success: bool) -> None:
        """Mark the first pending step that matches *tool_name* as done (if successful)."""
        if not success:
            return
        for step in self.plan.steps:
            if step.id in self._done:
                continue
            if step.tool and step.tool == tool_name:
                self._done[step.id] = True
                step.status = "done"
                return
        # No exact tool match — mark the first pending non-tool step if any
        for step in self.plan.steps:
            if step.id in self._done:
                continue
            if step.tool is None:
                self._done[step.id] = True
                step.status = "done"
                return

    def get_progress_prompt(self) -> str:
        """Return a <plan_progress> block showing done/remaining steps."""
        lines = ["<plan_progress>"]
        for step in self.plan.steps:
            marker = "✓" if step.id in self._done else "○"
            tool_hint = f" [{step.tool}]" if step.tool else ""
            lines.append(f"  {marker} {step.id}. {step.description}{tool_hint}")
        done = len(self._done)
        total = len(self.plan.steps)
        lines.append(f"\nProgress: {done}/{total} steps complete")
        lines.append("</plan_progress>")
        return "\n".join(lines)

    @property
    def is_complete(self) -> bool:
        """True when all plan steps are marked done."""
        return len(self._done) >= len(self.plan.steps)

    def get_current_step(self) -> Optional[PlanStep]:
        """Return the first pending step, or None if all done."""
        for step in self.plan.steps:
            if step.id not in self._done:
                return step
        return None
