"""Unified task template system.

Merges the SKILL.md skill system and browser YAML skill system into a
single discoverable, variable-substitutable template layer.

Templates can be:
  - SKILL.md files (markdown with YAML frontmatter)
  - Browser skill YAMLs (domain-specific procedures)
  - Task template YAMLs (new format with inputs + steps)

All three are surfaced through a single ``TemplateLibrary`` so the
agent can discover and apply the right template regardless of format.
"""

from __future__ import annotations

import re
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

from loguru import logger

try:
    import yaml
except ImportError:
    yaml = None  # type: ignore[assignment]


# ---------------------------------------------------------------------------
# Data model
# ---------------------------------------------------------------------------

@dataclass(frozen=True)
class TemplateInput:
    """A named input variable for a task template."""
    name: str
    description: str = ""
    default: str = ""


@dataclass(frozen=True)
class TaskTemplate:
    """Unified representation of any template type."""
    name: str
    description: str
    source: str          # "workspace" | "builtin" | "browser_skill"
    format: str          # "skill_md" | "browser_yaml" | "task_yaml"
    path: str            # absolute path to the source file
    content: str         # raw content (markdown body or formatted YAML)
    inputs: list[TemplateInput] = field(default_factory=list)
    role_hint: str = ""  # suggested agent role (e.g. "browser_operator")
    tools_needed: list[str] = field(default_factory=list)
    keywords: list[str] = field(default_factory=list)  # for auto-matching


# ---------------------------------------------------------------------------
# Variable substitution
# ---------------------------------------------------------------------------

_VAR_PATTERN = re.compile(r"\{\{(\w+)\}\}")


def substitute_variables(content: str, variables: dict[str, str]) -> str:
    """Replace ``{{name}}`` placeholders with provided values.

    Unknown variables are left as-is so the agent sees them as unfilled.
    """
    def _replace(match: re.Match) -> str:
        key = match.group(1)
        return variables.get(key, match.group(0))

    return _VAR_PATTERN.sub(_replace, content)


def extract_variable_names(content: str) -> list[str]:
    """Return unique variable names found in ``{{name}}`` placeholders."""
    return sorted(set(_VAR_PATTERN.findall(content)))


# ---------------------------------------------------------------------------
# Template parsing
# ---------------------------------------------------------------------------

_FRONTMATTER_RE = re.compile(r"^---\n(.*?)\n---\n?", re.DOTALL)


def _parse_frontmatter(text: str) -> tuple[dict[str, str], str]:
    """Split YAML frontmatter from markdown body. Returns (meta, body)."""
    match = _FRONTMATTER_RE.match(text)
    if not match:
        return {}, text

    meta: dict[str, str] = {}
    for line in match.group(1).split("\n"):
        if ":" in line:
            key, value = line.split(":", 1)
            meta[key.strip()] = value.strip().strip("\"'")
    return meta, text[match.end():].strip()


def _parse_skill_md(path: Path, source: str) -> TaskTemplate | None:
    """Parse a SKILL.md file into a TaskTemplate."""
    try:
        raw = path.read_text(encoding="utf-8")
    except Exception as exc:
        logger.warning("Failed to read skill {}: {}", path, exc)
        return None

    meta, body = _parse_frontmatter(raw)
    name = meta.get("name", path.parent.name)
    description = meta.get("description", "")
    role_hint = meta.get("role_hint", meta.get("role", ""))
    tools_raw = meta.get("tools_needed", "")
    tools_needed = [t.strip() for t in tools_raw.split(",") if t.strip()] if tools_raw else []

    # Extract inputs from frontmatter or auto-detect from {{vars}}
    inputs = _extract_inputs_from_meta(meta)
    if not inputs:
        var_names = extract_variable_names(body)
        inputs = [TemplateInput(name=v) for v in var_names]

    # Keywords for auto-matching: name tokens + description tokens
    keywords = _tokenize_keywords(name, description)

    return TaskTemplate(
        name=name,
        description=description or name,
        source=source,
        format="skill_md",
        path=str(path),
        content=body,
        inputs=inputs,
        role_hint=role_hint,
        tools_needed=tools_needed,
        keywords=keywords,
    )


def _parse_task_yaml(path: Path, source: str) -> TaskTemplate | None:
    """Parse a task template YAML (new format with inputs + steps)."""
    if yaml is None:
        return None
    try:
        data = yaml.safe_load(path.read_text(encoding="utf-8"))
    except Exception as exc:
        logger.warning("Failed to read task template {}: {}", path, exc)
        return None

    if not isinstance(data, dict):
        return None

    name = data.get("name", path.stem)
    description = data.get("description", "")
    role_hint = data.get("role_hint", "")
    tools_needed = data.get("tools_needed", [])

    # Parse inputs
    inputs = []
    for inp in data.get("inputs", []):
        if isinstance(inp, str):
            inputs.append(TemplateInput(name=inp))
        elif isinstance(inp, dict):
            inputs.append(TemplateInput(
                name=inp.get("name", ""),
                description=inp.get("description", ""),
                default=inp.get("default", ""),
            ))

    # Build content from steps
    steps = data.get("steps", [])
    content_parts = []
    if description:
        content_parts.append(f"## {name}\n\n{description}\n")
    if steps:
        content_parts.append("## Steps\n")
        for i, step in enumerate(steps, 1):
            content_parts.append(f"{i}. {step}")

    content = "\n".join(content_parts)
    keywords = _tokenize_keywords(name, description)

    return TaskTemplate(
        name=name,
        description=description or name,
        source=source,
        format="task_yaml",
        path=str(path),
        content=content,
        inputs=inputs,
        role_hint=role_hint,
        tools_needed=tools_needed or [],
        keywords=keywords,
    )


def _extract_inputs_from_meta(meta: dict[str, str]) -> list[TemplateInput]:
    """Extract input definitions from frontmatter string."""
    raw = meta.get("inputs", "")
    if not raw:
        return []
    # Simple comma-separated names
    return [TemplateInput(name=n.strip()) for n in raw.split(",") if n.strip()]


def _tokenize_keywords(name: str, description: str) -> list[str]:
    """Extract lowercase keyword tokens from name and description."""
    text = f"{name} {description}".lower()
    # Split on non-alphanumeric, deduplicate
    tokens = re.findall(r"[a-z0-9]+", text)
    # Filter stopwords
    stopwords = {"a", "an", "the", "to", "for", "and", "or", "in", "on", "of", "is", "it", "with", "use", "when"}
    return sorted(set(t for t in tokens if t not in stopwords and len(t) > 2))


# ---------------------------------------------------------------------------
# Template library
# ---------------------------------------------------------------------------

class TemplateLibrary:
    """Discovers, loads, and matches task templates from all sources."""

    def __init__(self, workspace: Path, builtin_skills_dir: Path | None = None):
        self._workspace = workspace
        self._builtin_skills = builtin_skills_dir or (Path(__file__).parent.parent / "skills")
        self._templates: dict[str, TaskTemplate] = {}
        self._loaded = False

    def load(self) -> None:
        """Discover and load all templates."""
        # 1. Workspace templates (highest priority)
        self._load_workspace_templates()

        # 2. Workspace skills (SKILL.md)
        self._load_skill_mds(self._workspace / "skills", "workspace")

        # 3. Built-in skills
        self._load_skill_mds(self._builtin_skills, "builtin")

        # 4. Browser skill YAMLs (via the existing skill_loader)
        self._load_browser_skills()

        self._loaded = True
        logger.info(
            "Template library loaded: {} templates ({} workspace, {} builtin, {} browser)",
            len(self._templates),
            sum(1 for t in self._templates.values() if t.source == "workspace"),
            sum(1 for t in self._templates.values() if t.source == "builtin"),
            sum(1 for t in self._templates.values() if t.source == "browser_skill"),
        )

    def _load_workspace_templates(self) -> None:
        """Load task template YAMLs from workspace/templates/."""
        templates_dir = self._workspace / "templates"
        if not templates_dir.exists():
            return
        for path in sorted(templates_dir.glob("*.yaml")):
            tpl = _parse_task_yaml(path, "workspace")
            if tpl and tpl.name not in self._templates:
                self._templates[tpl.name] = tpl
        for path in sorted(templates_dir.glob("*.yml")):
            tpl = _parse_task_yaml(path, "workspace")
            if tpl and tpl.name not in self._templates:
                self._templates[tpl.name] = tpl

    def _load_skill_mds(self, skills_dir: Path, source: str) -> None:
        """Load SKILL.md files from a skills directory."""
        if not skills_dir or not skills_dir.exists():
            return
        for skill_dir in sorted(skills_dir.iterdir()):
            if not skill_dir.is_dir():
                continue
            skill_file = skill_dir / "SKILL.md"
            if skill_file.exists() and skill_dir.name not in self._templates:
                tpl = _parse_skill_md(skill_file, source)
                if tpl:
                    self._templates[tpl.name] = tpl

    def _load_browser_skills(self) -> None:
        """Wrap browser skill YAMLs as templates via the existing SkillLibrary."""
        from madpod_agent.agent.browser.skill_loader import get_skill_library

        library = get_skill_library()
        for domain in library.get_all_domains():
            skill = library.get(domain)
            if skill is None or skill.domain in self._templates:
                continue
            self._templates[skill.domain] = TaskTemplate(
                name=skill.domain,
                description=f"Browser automation procedures for {skill.domain}",
                source="browser_skill",
                format="browser_yaml",
                path="",  # loaded via SkillLibrary, no single file path
                content=skill.format_for_llm(),
                role_hint="browser_operator",
                tools_needed=["browser_navigate", "browser_click", "browser_type"],
                keywords=_tokenize_keywords(skill.domain, " ".join(skill.procedures.keys())),
            )

    def get(self, name: str) -> TaskTemplate | None:
        """Get a template by exact name."""
        if not self._loaded:
            self.load()
        return self._templates.get(name)

    def list_all(self) -> list[TaskTemplate]:
        """Return all loaded templates."""
        if not self._loaded:
            self.load()
        return list(self._templates.values())

    def match(self, task_description: str, max_results: int = 3) -> list[TaskTemplate]:
        """Find templates matching a task description by keyword overlap.

        Returns up to ``max_results`` templates sorted by relevance score.
        """
        if not self._loaded:
            self.load()

        query_tokens = set(re.findall(r"[a-z0-9]+", task_description.lower()))
        scored: list[tuple[float, TaskTemplate]] = []

        for tpl in self._templates.values():
            if not tpl.keywords:
                continue
            tpl_keywords = set(tpl.keywords)
            overlap = query_tokens & tpl_keywords
            if not overlap:
                continue
            # Jaccard-like score weighted toward template coverage
            score = len(overlap) / len(tpl_keywords)
            scored.append((score, tpl))

        scored.sort(key=lambda x: x[0], reverse=True)
        return [tpl for _, tpl in scored[:max_results]]

    def format_for_context(self, templates: list[TaskTemplate], variables: dict[str, str] | None = None) -> str:
        """Format matched templates for injection into the agent's context."""
        if not templates:
            return ""

        parts = ["<matched_templates>"]
        for tpl in templates:
            content = tpl.content
            if variables:
                content = substitute_variables(content, variables)

            parts.append(f"\n### Template: {tpl.name}")
            if tpl.description:
                parts.append(f"Description: {tpl.description}")
            if tpl.inputs:
                unfilled = [inp for inp in tpl.inputs if not (variables or {}).get(inp.name)]
                if unfilled:
                    parts.append("Required inputs: " + ", ".join(
                        f"{inp.name} ({inp.description})" if inp.description else inp.name
                        for inp in unfilled
                    ))
            if tpl.role_hint:
                parts.append(f"Suggested role: {tpl.role_hint}")
            parts.append("")
            parts.append(content)

        parts.append("\n</matched_templates>")
        return "\n".join(parts)

    def build_summary(self) -> str:
        """Build compact XML summary of all templates for system prompt."""
        if not self._loaded:
            self.load()

        if not self._templates:
            return ""

        lines = ["<available_templates>"]
        for tpl in sorted(self._templates.values(), key=lambda t: t.name):
            source_tag = f' source="{tpl.source}"'
            inputs_tag = ""
            if tpl.inputs:
                inputs_str = ", ".join(inp.name for inp in tpl.inputs)
                inputs_tag = f' inputs="{inputs_str}"'
            lines.append(
                f'  <template name="{tpl.name}"{source_tag}{inputs_tag}>'
                f"{tpl.description}"
                f"</template>"
            )
        lines.append("</available_templates>")
        return "\n".join(lines)
