# Copyright (c) 2024-2025 Alain Prasquier - Supervaize.com. All rights reserved.
#
# This Source Code Form is subject to the terms of the Mozilla Public License, v. 2.0.
# If a copy of the MPL was not distributed with this file, you can obtain one at
# https://mozilla.org/MPL/2.0/.


from typing import TYPE_CHECKING, Any, ClassVar, Dict, Optional, Union, cast

import httpx
from pydantic import ConfigDict, Field, field_validator

from supervaizer.__version__ import VERSION
from supervaizer.common import ApiError, ApiResult, ApiSuccess, SvBaseModel
from supervaizer.telemetry import Telemetry

if TYPE_CHECKING:
    from supervaizer.agent import Agent
    from supervaizer.case import Case, CaseNodeUpdate
    from supervaizer.event import Event
    from supervaizer.job import Job
    from supervaizer.server import Server
from supervaizer.common import log


class AccountAbstract(SvBaseModel):
    """
    Account model for the Supervaize Control API.

    This represents an account that can be used to authenticate and communicate
    with the Supervaize SaaS API. Account parameters are provided by Supervaize.com.
    Create an account at [https://supervaize.com/](https://supervaize.com/)

    The account ID is generated by Supervaize.com. The account API key is provided
    by Supervaize.com and is used to authenticate requests to the Supervaize Control API.
    The API URL is provided by Supervaize.com and is the URL of the Supervaize SaaS API.

    The account provides methods for registering servers and agents, sending events,
    and communicating with the Supervaize platform.



    Attributes:
        workspace_id (str): The workspace ID provided by Supervaize.com
        api_key (str): The API key provided by Supervaize.com for authentication
        api_url (str): The URL of the Supervaize SaaS API provided by Supervaize.com
    """

    supervaizer_VERSION: ClassVar[str] = VERSION
    workspace_id: str = Field(description="The workspace ID provided by Supervaize.com")
    api_key: str = Field(
        description="The API key provided by Supervaize.com for authentication"
    )
    api_url: str = Field(
        description="The URL of the Supervaize SaaS API provided by Supervaize.com"
    )

    @field_validator("workspace_id", "api_key", "api_url", mode="before")
    @classmethod
    def strip_whitespace(cls, v: Any) -> Any:
        """Strip leading/trailing whitespace (e.g. newlines from env vars)."""
        if isinstance(v, str):
            return v.strip()
        return v

    model_config = cast(
        ConfigDict,
        {
            "reference_group": "Core",
            "json_schema_extra": {
                "examples": [
                    {
                        "workspace_id": "ws_1234567890abcdef",
                        "api_key": "sk_1234567890abcdef",
                        "api_url": "https://app.supervaize.com",
                    }
                ]
            },
        },
    )


class Account(AccountAbstract):
    # URL patterns dictionary
    _URL_PATTERNS = {
        "team": "{api_url}/w/{workspace_id}",
        "event": "{api_url}/w/{workspace_id}/api/v1/ctrl-events/",
        "agent_by_id": "{api_url}/w/{workspace_id}/api/v1/agents/{agent_id}",
        "agent_by_slug": "{api_url}/w/{workspace_id}/api/v1/agents/by-slug/{agent_slug}",
        "telemetry": "{api_url}/{telemetry_version}/telemetry",
    }

    def __init__(self, **kwargs: Any) -> None:
        """Account parameters are provided by Supervaize.com.
        The account ID is generated by Supervaize.com. The account API key is provided by Supervaize.com and is used
        to authenticate requests to the Supervaize Control API.
        The API URL is provided by Supervaize.com and is the URL of the Supervaize
        SaaS API.

        Raises:
            ValueError: If the account ID does not match the account name

        Tested in tests/test_account.py
        """
        super().__init__(**kwargs)

    @property
    def api_url_w_v1(self) -> str:
        """URL for the Supervaize Control API.
        Tested in tests/test_account.py
        """
        return f"{self.api_url}/w/{self.workspace_id}/api/v1"

    @property
    def api_headers(self) -> Dict[str, str]:
        """Headers for the Supervaize Control API.
        Tested in tests/test_account.py
        """
        return {
            "Authorization": f"Api-Key {self.api_key}",
            "accept": "application/json",
            "workspace": self.workspace_id,
        }

    @property
    def api_url_team(self) -> str:
        """URL for the Supervaize workspace team."""
        return self.get_url("team")

    @property
    def url_event(self) -> str:
        """URL for the Supervaize Control API.
        Tested in tests/test_account.py
        """
        return f"{self.api_url_w_v1}/ctrl-events/".strip()

    def get_url(self, pattern_name: str, **kwargs: Any) -> str:
        """Generate a URL using the predefined patterns.

        Args:
            pattern_name (str): The name of the URL pattern to use
            **kwargs: Additional parameters for URL formatting

        Returns:
            str: The formatted URL

        Raises:
            KeyError: If the pattern_name is not found in _URL_PATTERNS
        """
        if pattern_name not in self._URL_PATTERNS:
            raise KeyError(f"URL pattern '{pattern_name}' not found")

        pattern = self._URL_PATTERNS[pattern_name]

        # Default values from the account
        url_params = {
            "api_url": self.api_url,
            "workspace_id": self.workspace_id,
            "telemetry_version": "v1",  # Default telemetry version
        }

        # Override with any provided kwargs
        url_params.update(kwargs)

        return pattern.format(**url_params)

    def send_event(
        self,
        sender: Union["Agent", "Job", "Server", "Case", "CaseNodeUpdate"],
        event: "Event",
    ) -> ApiResult:
        """Send an event to the Supervaize SaaS API.

        Args:
            sender (Union[Agent, Server, Case, CaseNodeUpdate]): The sender of the event
            event (Event): The event to be sent

        Returns:
            ApiResult: ApiSuccess with response details if successful,
        Raises:
            Request exception if the request fails.
        """
        # Import here to avoid circular imports
        from supervaizer.account_service import send_event as service_send_event

        return service_send_event(self, sender, event)

    def register_server(self, server: "Server") -> ApiResult:
        """Register a server with the Supervaize Control API.

        Args:
            server (Server): The server to register.

        Returns:
            ApiResult: ApiSuccess with response details if successful,
                      ApiError with error details if request fails

        Side effects:
            - Sends a ServerRegisterEvent to the Supervaize Control API
        """
        # Import here to avoid circular imports
        from supervaizer.event import ServerRegisterEvent

        event = ServerRegisterEvent(server=server, account=self)
        result = self.send_event(sender=server, event=event)
        if isinstance(result, ApiSuccess):
            log.success(result.message)
            # TODO: Update server with the server ID from the response. store this ID in env variable.
        else:
            log.error(result.message)
        return result

    def _create_api_result(
        self,
        success: bool,
        message: str,
        detail: Any = None,
        url: str = "",
        exception: Optional[Exception] = None,
    ) -> ApiResult:
        if success:
            return ApiSuccess(message=message, detail=detail)
        return ApiError(message=message, url=url, payload=None, exception=exception)

    def get_agent_by(
        self, agent_id: str | None = None, agent_slug: str | None = None
    ) -> ApiSuccess | ApiError:
        """Get an agent from the Supervaize Control API.

        Args:
            agent_id (str): The ID of the agent to get.
            agent_slut (str): The name of the agent to get. (unique per workspace)

        Returns:
            ApiResult: ApiSuccess with response details if successful,ApiError with error details if request fails

        """
        url = ""
        if agent_id:
            url = f"{self.api_url_w_v1}/agents/{agent_id}"
        elif agent_slug:
            url = f"{self.api_url_w_v1}/agents/by-slug/{agent_slug}"
        else:
            raise ValueError("No agent ID or slug provided")

        headers = self.api_headers

        try:
            response = httpx.get(url, headers=headers, follow_redirects=True)
            response.raise_for_status()
            return ApiSuccess(
                message=f"GET Agent <{agent_slug or agent_id}>",
                detail=response.json(),
            )
        except httpx.HTTPError as e:
            return ApiError(
                message=f"Error GET Agent <{agent_slug or agent_id}>",
                url=url,
                payload=None,
                exception=e,
            )

    def register_agent(self, agent: "Agent", polling: bool = True) -> ApiResult:
        """Send a registration event to the Supervaize Control API.
            This will be used for polling, when the agent is registered without a server.
        Args:
            agent (Agent): The agent sending the registration event
            polling (bool): If server is not defined, polling will be used.

        Returns:
            ApiResult: ApiSuccess with response details if successful,
                      ApiError with error details if request fails
        """
        # Import here to avoid circular imports
        from supervaizer.event import AgentRegisterEvent

        event = AgentRegisterEvent(agent=agent, account=self, polling=polling)
        return self.send_event(agent, event)

    def send_start_case(self, case: "Case") -> ApiResult:
        # Import here to avoid circular imports
        from supervaizer.event import CaseStartEvent

        event = CaseStartEvent(case=case, account=self)
        return self.send_event(case, event)

    def send_update_case(self, case: "Case", update: "CaseNodeUpdate") -> ApiResult:
        # Import here to avoid circular imports
        log.debug(f"[send_update_case] CaseRef {case} with update {update}")
        log.debug(f"[send_update_case] {type(case)}")
        log.debug(f"[send_update_case] {type(update)}")
        from supervaizer.event import CaseUpdateEvent

        event = CaseUpdateEvent(case=case, update=update, account=self)
        return self.send_event(update, event)

    def send_telemetry(self, telemetry: Telemetry) -> ApiResult:
        """Send telemetry data to the Supervaize Control API.

        Args:
            telemetry (Telemetry): The telemetry object to be sent

        Returns:
            ApiResult: ApiSuccess with response details if successful,
                      ApiError with error details if request fails
        """
        url = self.get_url("telemetry")
        headers = self.api_headers
        payload = {"workspace_id": self.workspace_id} | telemetry.payload
        result: ApiSuccess | ApiError
        try:
            response = httpx.post(url, headers=headers, json=payload)
            response.raise_for_status()
            result = ApiSuccess(
                message=f"Telemetry sent {telemetry.type.name}", detail=response.text
            )
        except httpx.HTTPError as e:
            result = ApiError(
                message=f"Error sending telemetry {telemetry.type.name}",
                url=url,
                payload=payload,
                exception=e,
            )

        return result
