import os
import time
import importlib
from typing import List, Dict, Optional, Callable, Any
from datetime import datetime
from .card import AgentCard


PROVIDER_MAP = {
    "groq":      ("synapse.providers.groq",      "GroqProvider"),
    "openai":    ("synapse.providers.openai",     "OpenAIProvider"),
    "anthropic": ("synapse.providers.anthropic",  "AnthropicProvider"),
    "gemini":    ("synapse.providers.gemini",     "GeminiProvider"),
}

DEFAULT_MODELS = {
    "groq":      "llama-3.3-70b-versatile",
    "openai":    "gpt-4o",
    "anthropic": "claude-sonnet-4-20250514",
    "gemini":    "gemini-1.5-flash",
}

API_KEY_ENV = {
    "groq":      "GROQ_API_KEY",
    "openai":    "OPENAI_API_KEY",
    "anthropic": "ANTHROPIC_API_KEY",
    "gemini":    "GEMINI_API_KEY",
}


class Agent:
    """
    Core Agent class.

    Features:
    1.  Multi-LLM support (Groq, OpenAI, Anthropic, Gemini)
    2.  Local agent-to-agent communication
    6.  AgentCard identity/profile
    10. Agent versioning
    11. Agent cloning
    12. Agent hot-reload
    13. Agent lifecycle hooks (on_start, on_stop, on_message)
    14. Agent health checks
    15. Agent heartbeat
    """

    def __init__(
        self,
        name: str,
        provider: str = "groq",
        model: Optional[str] = None,
        api_key: Optional[str] = None,
        system_prompt: Optional[str] = None,
        skills: Optional[List[str]] = None,
        tools: Optional[List[Any]] = None,
        version: str = "1.0.0",
        metadata: Optional[dict] = None,
        on_start: Optional[Callable] = None,
        on_stop: Optional[Callable] = None,
        on_message: Optional[Callable] = None,
        max_retries: int = 3,
        timeout: int = 30,
        verbose: bool = False,
    ):
        self.name = name
        self.provider_name = provider.lower()
        self.model = model or DEFAULT_MODELS.get(self.provider_name, "")
        self.system_prompt = system_prompt or f"You are {name}, a helpful AI agent in the Synapse multi-agent system."
        self.skills = skills or []
        self.tools = tools or []
        self.version = version
        self.metadata = metadata or {}
        self.max_retries = max_retries
        self.timeout = timeout
        self.verbose = verbose
        self._status = "initialized"
        self._message_count = 0
        self._total_tokens = 0
        self._total_cost = 0.0
        self._created_at = datetime.utcnow().isoformat()
        self._last_heartbeat = None

        # Lifecycle hooks
        self._on_start = on_start
        self._on_stop = on_stop
        self._on_message = on_message

        # Resolve API key
        env_key = API_KEY_ENV.get(self.provider_name, "")
        self.api_key = api_key or os.getenv(env_key, "")
        if not self.api_key:
            raise ValueError(
                f"No API key for provider '{self.provider_name}'. "
                f"Pass api_key= or set {env_key} environment variable."
            )

        self._provider = self._load_provider()
        self._tool_instances = {t.name: t for t in self.tools} if self.tools else {}

        # Feature 13: trigger on_start hook
        if self._on_start:
            self._on_start(self)
        self._status = "ready"

    def _load_provider(self):
        if self.provider_name not in PROVIDER_MAP:
            raise ValueError(
                f"Unknown provider '{self.provider_name}'. "
                f"Choose from: {list(PROVIDER_MAP.keys())}"
            )
        module_path, class_name = PROVIDER_MAP[self.provider_name]
        module = importlib.import_module(module_path)
        cls = getattr(module, class_name)
        return cls(api_key=self.api_key, model=self.model)

    def complete(self, messages: List[Dict], use_tools: bool = True) -> str:
        """Send messages and get a response. Feature 17: retries on failure."""
        # Feature 13: on_message hook
        if self._on_message:
            self._on_message(self, messages)

        last_error = None
        for attempt in range(self.max_retries):
            try:
                if self.verbose:
                    print(f"  [{self.name}] attempt {attempt + 1}/{self.max_retries}")

                # Feature 41-60: use tools if available
                if use_tools and self._tool_instances:
                    response = self._complete_with_tools(messages)
                else:
                    response = self._provider.complete(messages, system_prompt=self.system_prompt)

                self._message_count += 1
                self._last_heartbeat = datetime.utcnow().isoformat()
                return response

            except Exception as e:
                last_error = e
                if attempt < self.max_retries - 1:
                    time.sleep(2 ** attempt)  # exponential backoff
                    if self.verbose:
                        print(f"  [{self.name}] retrying after error: {e}")

        raise RuntimeError(f"Agent '{self.name}' failed after {self.max_retries} attempts. Last error: {last_error}")

    def _complete_with_tools(self, messages: List[Dict]) -> str:
        """Run agent with tool support."""
        tool_descriptions = "\n".join([
            f"- {t.name}: {t.description}" for t in self.tools
        ])
        tool_system = (
            f"{self.system_prompt}\n\n"
            f"You have access to these tools:\n{tool_descriptions}\n\n"
            f"To use a tool, write: TOOL: <tool_name> | <input>"
        )

        response = self._provider.complete(messages, system_prompt=tool_system)

        # Parse tool calls
        if "TOOL:" in response:
            lines = response.split("\n")
            final_parts = []
            for line in lines:
                if line.startswith("TOOL:"):
                    parts = line.replace("TOOL:", "").split("|")
                    if len(parts) == 2:
                        tool_name = parts[0].strip()
                        tool_input = parts[1].strip()
                        if tool_name in self._tool_instances:
                            tool_result = self._tool_instances[tool_name].run(tool_input)
                            final_parts.append(f"[{tool_name} result]: {tool_result}")
                        else:
                            final_parts.append(f"[Tool {tool_name} not found]")
                else:
                    final_parts.append(line)

            # Second pass: use tool results to get final answer
            tool_context = "\n".join(final_parts)
            followup = messages + [
                {"role": "assistant", "content": response},
                {"role": "user", "content": f"Tool results:\n{tool_context}\n\nNow give your final answer."}
            ]
            return self._provider.complete(followup, system_prompt=self.system_prompt)

        return response

    # Feature 14: health check
    def health(self) -> dict:
        return {
            "name": self.name,
            "status": self._status,
            "provider": self.provider_name,
            "model": self.model,
            "message_count": self._message_count,
            "last_heartbeat": self._last_heartbeat,
            "created_at": self._created_at,
            "tools": list(self._tool_instances.keys()),
        }

    # Feature 15: heartbeat
    def heartbeat(self) -> str:
        self._last_heartbeat = datetime.utcnow().isoformat()
        return self._last_heartbeat

    # Feature 10: versioning
    def get_version(self) -> str:
        return self.version

    # Feature 11: clone
    def clone(self, new_name: str) -> "Agent":
        return Agent(
            name=new_name,
            provider=self.provider_name,
            model=self.model,
            api_key=self.api_key,
            system_prompt=self.system_prompt,
            skills=self.skills.copy(),
            tools=self.tools.copy(),
            version=self.version,
            metadata=self.metadata.copy(),
            max_retries=self.max_retries,
            timeout=self.timeout,
            verbose=self.verbose,
        )

    # Feature 12: hot-reload system prompt
    def reload(self, system_prompt: str = None, model: str = None):
        if system_prompt:
            self.system_prompt = system_prompt
        if model:
            self.model = model
            self._provider = self._load_provider()
        if self.verbose:
            print(f"  [{self.name}] reloaded.")

    # Feature 6: get agent card
    def card(self) -> AgentCard:
        return AgentCard(
            name=self.name,
            description=self.system_prompt[:100],
            skills=self.skills,
            tools=list(self._tool_instances.keys()),
            provider=self.provider_name,
            model=self.model,
            version=self.version,
            metadata=self.metadata,
        )

    def stop(self):
        self._status = "stopped"
        if self._on_stop:
            self._on_stop(self)

    def __repr__(self):
        return f"Agent(name={self.name!r}, provider={self.provider_name!r}, model={self.model!r}, tools={list(self._tool_instances.keys())})"
