import logging
import uuid
from datetime import datetime, timezone
from typing import Any, Optional
from .models import AgentModel, AgentVersion, AgentRun, Event
from .client import EzopClient

logger = logging.getLogger(__name__)


class Span:
    """Context manager that wraps a scoped duration as a pair of events.

    On enter: emits an ``in_progress`` event.
    On exit:  emits an ``ok`` or ``error`` event with the same ``span_id``.

    The ``span_id`` is shared between both events so duration can be
    reconstructed by grouping on ``span_id`` and computing the time delta.

    Nested spans automatically propagate context via ``_current_span_id`` on
    the active run, so each span records its ``parent_id`` correctly.
    """

    def __init__(
        self,
        agent: "Agent",
        name: str,
        category: str,
        *,
        input: Optional[Any] = None,
        metadata: Optional[dict] = None,
    ):
        self._agent = agent
        self.name = name
        self.category = category
        self.span_id = str(uuid.uuid4())
        self._input = input
        self._metadata = metadata
        self._output: Optional[Any] = None
        self._parent_id: Optional[str] = None
        self._run: Optional[AgentRun] = None

    def set_output(self, output: Any) -> None:
        """Capture the output to be sent with the closing event."""
        self._output = output

    def __enter__(self) -> "Span":
        if self._agent.current_run is None:
            raise RuntimeError("No active run. Call start_run() first.")

        # Snapshot the run so this span is bound to it for its entire lifetime
        self._run = self._agent.current_run

        # Capture the enclosing span before overwriting it
        self._parent_id = self._run._current_span_id
        self._run._current_span_id = self.span_id

        self._agent.emit(
            name=self.name,
            category=self.category,
            span_id=self.span_id,
            parent_id=self._parent_id,
            status="in_progress",
            input=self._input,
            metadata=self._metadata,
        )
        logger.debug(
            "Span entered name=%s span_id=%s parent_id=%s",
            self.name,
            self.span_id,
            self._parent_id,
        )
        return self

    def __exit__(self, exc_type, exc_val, _) -> bool:
        status = "error" if exc_type else "ok"
        error = str(exc_val) if exc_val else None

        self._agent.emit(
            name=self.name,
            category=self.category,
            span_id=self.span_id,
            parent_id=self._parent_id,
            status=status,
            output=self._output,
            metadata=self._metadata,
            error=error,
        )
        logger.debug(
            "Span exited name=%s span_id=%s status=%s", self.name, self.span_id, status
        )

        # Restore the enclosing span on the snapshotted run
        self._run._current_span_id = self._parent_id
        return False  # never suppress exceptions


class Agent:

    def __init__(self, model: AgentModel, version: AgentVersion):
        self.model = model
        self.version = version
        self.client = EzopClient()
        self.current_run: Optional[AgentRun] = None

    @classmethod
    def init(
        cls,
        name: str,
        owner: str,
        version: str,
        runtime: str,
        *,
        description: Optional[str] = None,
        default_permissions: Optional[list[str]] = None,
        permissions: Optional[list[str]] = None,
        changelog: Optional[str] = None,
    ) -> "Agent":
        """
        Initialize an Ezop agent and register it with the Ezop platform.
        Creates the agent record and a corresponding version record.
        """
        logger.info(
            "Initializing agent name=%s owner=%s version=%s runtime=%s",
            name,
            owner,
            version,
            runtime,
        )
        client = EzopClient()

        agent_data = client.register_agent(
            {
                "name": name,
                "owner": owner,
                "runtime": runtime,
                "description": description,
                "default_permissions": default_permissions,
            }
        )["data"]

        model = AgentModel(
            id=agent_data["id"],
            name=agent_data["name"],
            owner=agent_data["owner"],
            runtime=agent_data.get("runtime"),
            description=agent_data.get("description"),
            default_permissions=agent_data.get("default_permissions") or [],
        )
        logger.debug("Agent model built id=%s", model.id)

        version_data = client.create_version(
            model.id,
            {
                "version": version,
                "permissions": permissions,
                "changelog": changelog,
            },
        )["data"]

        agent_version = AgentVersion(
            id=version_data["id"],
            agent_id=model.id,
            version=version_data["version"],
            permissions=version_data.get("permissions") or [],
            changelog=version_data.get("changelog"),
        )
        logger.debug("Agent version built id=%s", agent_version.id)

        agent = cls(model, agent_version)

        logger.info(
            "Starting run for agent id=%s version_id=%s", model.id, agent_version.id
        )
        run_data = client.start_run(model.id, agent_version.id)["data"]
        agent.current_run = AgentRun(
            id=run_data["id"],
            agent_id=model.id,
            version_id=agent_version.id,
            status=run_data.get("status", "running"),
        )
        logger.info(
            "Agent initialized id=%s version_id=%s run_id=%s",
            model.id,
            agent_version.id,
            agent.current_run.id,
        )
        return agent

    def close(
        self,
        status: str = "running",
        total_tokens: Optional[int] = None,
        total_cost: Optional[float] = None,
        metadata: Optional[dict] = None,
        message: Optional[str] = None,
    ) -> AgentRun:
        """Close the current run and record its outcome.

        Args:
            status: Final status of the run (e.g. "success", "failed", "partial", "cancelled").
            message: Optional human-readable message, e.g. a failure reason.
        """
        if self.current_run is None:
            logger.error("close called with no active run")
            raise RuntimeError("No active run.")

        logger.info("Ending run id=%s status=%s", self.current_run.id, status)
        run_data = self.client.end_run(
            self.current_run.id,
            status=status,
            total_tokens=total_tokens,
            total_cost=total_cost,
            metadata=metadata,
            message=message,
        )["data"]
        self.current_run.status = run_data.get("status", status)
        self.current_run.total_tokens = run_data.get("total_tokens")
        self.current_run.total_cost = run_data.get("total_cost")
        self.current_run.metadata = run_data.get("metadata")
        self.current_run.message = run_data.get("message")
        logger.debug(
            "Run ended id=%s status=%s total_tokens=%s total_cost=%s",
            self.current_run.id,
            self.current_run.status,
            self.current_run.total_tokens,
            self.current_run.total_cost,
        )
        return self.current_run

    def emit(
        self,
        name: str,
        category: str,
        *,
        span_id: Optional[str] = None,
        parent_id: Optional[str] = None,
        status: Optional[str] = None,
        input: Optional[Any] = None,
        output: Optional[Any] = None,
        metadata: Optional[dict] = None,
        error: Optional[str] = None,
    ) -> Event:
        """Emit a single point-in-time event on the current run.

        If called inside a ``span``, the event is automatically parented to
        that span via ``parent_id``.
        """
        if self.current_run is None:
            raise RuntimeError("No active run. Call start_run() first.")

        resolved_parent_id = (
            parent_id if parent_id is not None else self.current_run._current_span_id
        )

        event = Event(
            id=str(uuid.uuid4()),
            run_id=self.current_run.id,
            span_id=span_id,
            parent_id=resolved_parent_id,
            name=name,
            category=category,
            status=status,
            timestamp=datetime.now(timezone.utc),
            input=input,
            output=output,
            metadata=metadata,
            error=error,
        )

        body: dict = {
            "id": event.id,
            "name": event.name,
            "category": event.category,
            "timestamp": event.timestamp.isoformat(),
        }
        if event.span_id is not None:
            body["span_id"] = event.span_id
        if event.parent_id is not None:
            body["parent_id"] = event.parent_id
        if event.status is not None:
            body["status"] = event.status
        if event.input is not None:
            body["input"] = event.input
        if event.output is not None:
            body["output"] = event.output
        if event.metadata is not None:
            body["metadata"] = event.metadata
        if event.error is not None:
            body["error"] = event.error

        self.client.emit_event(self.current_run.id, body)
        logger.debug(
            "Event emitted id=%s name=%s span_id=%s parent_id=%s",
            event.id,
            event.name,
            event.span_id,
            event.parent_id,
        )
        return event

    def span(
        self,
        name: str,
        category: str,
        *,
        input: Optional[Any] = None,
        metadata: Optional[dict] = None,
    ) -> Span:
        """Return a context manager that tracks a scoped duration as a span.

        Usage::

            with agent.span("llm.call", category="llm", input=prompt) as s:
                result = llm.generate(prompt)
                s.set_output(result)
        """
        return Span(self, name, category, input=input, metadata=metadata)
