# Copyright (c) Metis. All rights reserved.

"""Agent manifest parsing for Mantis platform.

Parses the mantis_manifest/ directory structure:
  mantis_manifest/
    agent.yml       — Agent identity, capabilities, resources, runner config
    task.yml        — (optional) Task input/output JSON Schema
    .env            — (optional, gitignored) API keys and secrets
"""

from __future__ import annotations

from dataclasses import asdict, dataclass, field
from pathlib import Path
from typing import Any, Dict, List, Optional, Union


@dataclass
class RunnerConfig:
    """Runner lifecycle configuration declared in agent.yml."""

    max_rollouts: Optional[int] = None  # None = unlimited
    poll_interval: float = 5.0  # seconds between store polls
    n_runners: int = 1  # parallel runners in one container
    container_lifecycle: str = "persistent"  # "persistent" or "ephemeral"


@dataclass
class ResourceConfig:
    """Container resource limits declared in agent.yml."""

    memory: Optional[str] = None  # Docker memory limit, e.g. "2g"
    cpu: Optional[float] = None  # CPU cores (fractional ok), e.g. 1.0
    gpu: Union[bool, int] = False  # False, True, or device count
    stop_timeout: int = 30  # graceful shutdown seconds


@dataclass
class AgentManifest:
    """Parsed agent.yml manifest."""

    name: str
    version: str = "1.0.0"
    capabilities: List[str] = field(default_factory=list)
    resources: Dict[str, Any] = field(default_factory=dict)
    description: str = ""
    runner: RunnerConfig = field(default_factory=RunnerConfig)
    resource_limits: ResourceConfig = field(default_factory=ResourceConfig)
    env: Dict[str, str] = field(default_factory=dict)


@dataclass
class TaskSchema:
    """Parsed task.yml schema."""

    input_schema: Dict[str, Any] = field(default_factory=dict)


def parse_manifest(manifest_dir: Path) -> AgentManifest:
    """Parse mantis_manifest/agent.yml into an AgentManifest.

    Args:
        manifest_dir: Path to the mantis_manifest/ directory.

    Returns:
        Parsed AgentManifest dataclass.

    Raises:
        FileNotFoundError: If agent.yml is not found.
        ValueError: If agent.yml is malformed.
    """
    try:
        import yaml
    except ImportError:
        raise ImportError("PyYAML is required for manifest parsing. Install with: pip install pyyaml")

    agent_yml = manifest_dir / "agent.yml"
    if not agent_yml.exists():
        raise FileNotFoundError(f"agent.yml not found in {manifest_dir}")

    with open(agent_yml) as f:
        data = yaml.safe_load(f)

    if not isinstance(data, dict):
        raise ValueError(f"agent.yml must be a YAML mapping, got {type(data).__name__}")

    if "name" not in data:
        raise ValueError("agent.yml must contain a 'name' field")

    # Parse runner config
    runner_raw = data.get("runner", {})
    runner = RunnerConfig(
        max_rollouts=runner_raw.get("max_rollouts"),
        poll_interval=runner_raw.get("poll_interval", 5.0),
        n_runners=runner_raw.get("n_runners", 1),
        container_lifecycle=runner_raw.get("container_lifecycle", "persistent"),
    )

    # Parse resource limits
    res_raw = data.get("resources", {})
    resource_limits = ResourceConfig(
        memory=res_raw.get("memory"),
        cpu=res_raw.get("cpu"),
        gpu=res_raw.get("gpu", False),
        stop_timeout=res_raw.get("stop_timeout", 30),
    )

    return AgentManifest(
        name=data["name"],
        version=data.get("version", "1.0.0"),
        capabilities=data.get("capabilities", []),
        resources=data.get("resources", {}),
        description=data.get("description", ""),
        runner=runner,
        resource_limits=resource_limits,
        env=data.get("env", {}),
    )


def parse_task_schema(manifest_dir: Path) -> Optional[TaskSchema]:
    """Parse mantis_manifest/task.yml into a TaskSchema.

    Args:
        manifest_dir: Path to the mantis_manifest/ directory.

    Returns:
        Parsed TaskSchema dataclass, or None if task.yml doesn't exist.
    """
    try:
        import yaml
    except ImportError:
        return None

    task_yml = manifest_dir / "task.yml"
    if not task_yml.exists():
        return None

    with open(task_yml) as f:
        data = yaml.safe_load(f)

    if not isinstance(data, dict):
        return None

    return TaskSchema(
        input_schema=data.get("input_schema", {}),
    )


def parse_secrets(manifest_dir: Path) -> Dict[str, str]:
    """Parse mantis_manifest/.env into a dict of secret name → value.

    Standard .env format: one KEY=VALUE per line, # comments, blank lines ignored.
    No external dependencies needed.

    Args:
        manifest_dir: Path to the mantis_manifest/ directory.

    Returns:
        Dict mapping secret names to values. Empty dict if .env doesn't exist.
    """
    env_file = manifest_dir / ".env"
    if not env_file.exists():
        return {}

    secrets: Dict[str, str] = {}
    with open(env_file) as f:
        for line in f:
            line = line.strip()
            # Skip blank lines and comments
            if not line or line.startswith("#"):
                continue
            # Split on first = only
            if "=" not in line:
                continue
            key, value = line.split("=", 1)
            key = key.strip()
            value = value.strip()
            # Strip surrounding quotes if present
            if len(value) >= 2 and value[0] == value[-1] and value[0] in ('"', "'"):
                value = value[1:-1]
            if key:
                secrets[key] = value

    return secrets
