from typing import List, Optional, Dict, Any
import urllib.request
import urllib.error
import json
import logging

from puli_models import (
    IncidentQueryResult,
    ChaosPatternQueryResult,
    BestPracticeQueryResult,
    SecurityQueryResult,
    RegulationQueryResult,
    DomainIndicatorQueryResult,
)

from .config import ProxyConfig
from .token_manager import get_machine_id, read_token, write_token, clear_token

logger = logging.getLogger(__name__)


class ProxyClient:
    """Client for interacting with the Puli proxy service."""

    def __init__(self, config: Optional[ProxyConfig] = None):
        """
        Initialize the proxy client.

        Args:
            config: ProxyConfig instance. If None, loads from environment variables.
        """
        self.config = config or ProxyConfig.from_env()
        self.role: Optional[str] = None
        self._ensure_authenticated()
        self.mcp_config = self._fetch_mcp_config()
        # Role is available from /config/mcp response (covers cached-token path too)
        if self.mcp_config.get("role"):
            self.role = self.mcp_config["role"]

    def _ensure_authenticated(self) -> None:
        """Load existing JWT or register to obtain one."""
        token = read_token(self.config.api_key)
        if token:
            self.config.jwt_token = token
            return
        self._register()

    def _register(self) -> None:
        """Register this machine and store the returned JWT."""
        machine_id = get_machine_id()
        url = f"{self.config.base_url}/auth/register"
        headers = {
            "X-API-Key": self.config.api_key,
            "Content-Type": "application/json",
        }
        if self.config.identity_token:
            headers["Authorization"] = f"Bearer {self.config.identity_token}"

        body = json.dumps({"machine_id": machine_id}).encode()
        req = urllib.request.Request(url, data=body, headers=headers, method="POST")

        try:
            with urllib.request.urlopen(req) as response:
                result = json.loads(response.read().decode())
        except urllib.error.HTTPError as e:
            error_body = e.read().decode() if e.fp else ""
            raise RuntimeError(f"Registration failed: {e.code} - {error_body}") from e

        token = result["token"]
        write_token(token, self.config.api_key)
        self.config.jwt_token = token
        self.role = result.get("role")

    def _fetch_mcp_config(self) -> Dict[str, Any]:
        """Fetch MCP configuration from the proxy after authentication."""
        try:
            return self._request("GET", "/config/mcp")
        except RuntimeError as e:
            # If config fetch fails, return empty config (allows local dev fallback)
            if "404" in str(e) or "500" in str(e):
                logger.warning("Failed to fetch MCP config from proxy: %s", e)
                logger.warning("Falling back to empty config. Client will use from_env() methods for local development.")
                return {"config": {}, "prompts": {}}
            raise

    def _request(
        self,
        method: str,
        path: str,
        data: Optional[dict] = None,
    ) -> dict:
        """Make an HTTP request to the proxy with auto-retry on 401."""
        url = f"{self.config.base_url}{path}"
        headers = {
            "Content-Type": "application/json",
        }

        if self.config.jwt_token:
            headers["X-Auth-Token"] = self.config.jwt_token

        if self.config.identity_token:
            headers["Authorization"] = f"Bearer {self.config.identity_token}"

        body = json.dumps(data).encode() if data else None
        request = urllib.request.Request(url, data=body, headers=headers, method=method)

        try:
            with urllib.request.urlopen(request) as response:
                return json.loads(response.read().decode())
        except urllib.error.HTTPError as e:
            if e.code == 401:
                # Token expired or invalid — re-register once and retry
                clear_token()
                self._register()
                return self._retry_request(method, url, data)
            error_body = e.read().decode() if e.fp else ""
            raise RuntimeError(f"Proxy request failed: {e.code} - {error_body}") from e

    def _retry_request(
        self,
        method: str,
        url: str,
        data: Optional[dict] = None,
    ) -> dict:
        """Single retry after re-registration (no infinite loop)."""
        headers = {
            "Content-Type": "application/json",
        }

        if self.config.jwt_token:
            headers["X-Auth-Token"] = self.config.jwt_token

        if self.config.identity_token:
            headers["Authorization"] = f"Bearer {self.config.identity_token}"

        body = json.dumps(data).encode() if data else None
        request = urllib.request.Request(url, data=body, headers=headers, method=method)

        try:
            with urllib.request.urlopen(request) as response:
                return json.loads(response.read().decode())
        except urllib.error.HTTPError as e:
            error_body = e.read().decode() if e.fp else ""
            raise RuntimeError(f"Proxy request failed after retry: {e.code} - {error_body}") from e

    def _search(
        self,
        query_vector: List[float],
        url: str,
        limit: int = 10,
        filter_expr: Optional[str] = None,
        output_fields: Optional[List[str]] = None,
    ) -> List[Any]:
        """
        Search for similar incidents by vector.

        Args:
            query_vector: The embedding vector to search with.
            limit: Maximum number of results to return.
            filter_expr: Optional filter expression.
            output_fields: Fields to include in results.

        Returns:
            List of matching IncidentQueryResult objects.
        """
        data = {
            "query_vector": query_vector,
            "limit": limit,
        }
        if filter_expr:
            data["filter"] = filter_expr
        if output_fields:
            data["output_fields"] = output_fields

        response = self._request("POST", url, data)
        return response.get("results", [])

    def search_incidents(
        self,
        query_vector: List[float],
        limit: int = 10,
        filter_expr: Optional[str] = None,
        output_fields: Optional[List[str]] = None,
    ) -> List[IncidentQueryResult]:
        results = self._search(query_vector, "/incidents/search", limit, filter_expr, output_fields)
        return [IncidentQueryResult.from_dict(r) for r in results]

    def search_chaos_patterns(
        self,
        query_vector: List[float],
        limit: int = 10,
        filter_expr: Optional[str] = None,
        output_fields: Optional[List[str]] = None,
    ) -> List[ChaosPatternQueryResult]:
        results = self._search(query_vector, "/chaos-patterns/search", limit, filter_expr, output_fields)
        return [ChaosPatternQueryResult.from_dict(r) for r in results]

    def search_best_practices(
        self,
        query_vector: List[float],
        limit: int = 10,
        filter_expr: Optional[str] = None,
        output_fields: Optional[List[str]] = None,
    ) -> List[BestPracticeQueryResult]:
        results = self._search(query_vector, "/best-practices/search", limit, filter_expr, output_fields)
        return [BestPracticeQueryResult.from_dict(r) for r in results]

    def search_security_items(
        self,
        query_vector: List[float],
        limit: int = 10,
        filter_expr: Optional[str] = None,
        output_fields: Optional[List[str]] = None,
    ) -> List[SecurityQueryResult]:
        results = self._search(query_vector, "/security/search", limit, filter_expr, output_fields)
        return [SecurityQueryResult.from_dict(r) for r in results]

    def search_regulations(
        self,
        query_vector: List[float],
        limit: int = 10,
        filter_expr: Optional[str] = None,
        output_fields: Optional[List[str]] = None,
    ) -> List[RegulationQueryResult]:
        results = self._search(query_vector, "/regulations/search", limit, filter_expr, output_fields)
        return [RegulationQueryResult.from_dict(r) for r in results]

    def search_domain_indicators(
        self,
        query_vector: List[float],
        limit: int = 10,
        filter_expr: Optional[str] = None,
        output_fields: Optional[List[str]] = None,
    ) -> List[DomainIndicatorQueryResult]:
        results = self._search(query_vector, "/domain-indicators/search", limit, filter_expr, output_fields)
        return [DomainIndicatorQueryResult.from_dict(r) for r in results]

    def query(
        self,
        filter_expr: str,
        output_fields: Optional[List[str]] = None,
        limit: int = 100,
    ) -> List[IncidentQueryResult]:
        """
        Query incidents by filter expression (no vector search).

        Args:
            filter_expr: Filter expression.
            output_fields: Fields to include in results.
            limit: Maximum number of results.

        Returns:
            List of matching IncidentQueryResult objects.
        """
        data = {
            "filter": filter_expr,
            "limit": limit,
        }
        if output_fields:
            data["output_fields"] = output_fields

        response = self._request("POST", "/incidents/query", data)
        return [IncidentQueryResult.from_dict(r) for r in response.get("results", [])]

    def insert(self, incidents: List[Dict[str, Any]]) -> Dict[str, Any]:
        """
        Insert incidents into the collection.

        Args:
            incidents: List of incident dictionaries to insert.

        Returns:
            Insert result with IDs of inserted records.
        """
        data = {"data": incidents}
        return self._request("POST", "/incidents/insert", data)

    def get_by_id(self, incident_id: str) -> Optional[Dict[str, Any]]:
        """
        Get a single incident by ID.

        Args:
            incident_id: The incident ID to retrieve.

        Returns:
            Incident dictionary or None if not found.
        """
        try:
            return self._request("GET", f"/incidents/{incident_id}")
        except RuntimeError as e:
            if "404" in str(e):
                return None
            raise

    def delete(self, ids: List[str]) -> Dict[str, Any]:
        """
        Delete incidents by IDs.

        Args:
            ids: List of incident IDs to delete.

        Returns:
            Delete result.
        """
        data = {"ids": ids}
        return self._request("DELETE", "/incidents", data)

    def count(self) -> int:
        """
        Get the total number of incidents in the collection.

        Returns:
            Number of incidents.
        """
        response = self._request("GET", "/incidents/count")
        return response.get("count", 0)

    def generate_embedding(self, text: str) -> List[float]:
        """
        Generate an embedding vector via the proxy's embedding endpoint.

        Args:
            text: The text to embed.

        Returns:
            A list of floats representing the embedding vector.
        """
        response = self._request("POST", "/embeddings/generate", {"text": text})
        return response["vector"]

    def get_tenant_profile(self, question: str) -> Dict[str, Any]:
        """Ask the tenant's RAGFlow profiler chat a question.

        The proxy resolves the profiler chat (``puli-{tenant_id}-profiler``),
        gets or creates a session, and runs completion — all in one call.

        Args:
            question: The prompt to send to the profiler chat.

        Returns:
            Dict with ``answer``, ``session_id``, ``chat_id``, ``tenant_id``.

        Raises:
            RuntimeError: If the proxy returns an error (no RAGFlow access,
                profiler chat not found, etc.).
        """
        return self._request("POST", "/ragflow/profiler/ask", {"question": question})

    def submit_report(self, session: Any) -> Dict[str, Any]:
        """Submit a coding session assessment summary to the proxy for storage.

        Args:
            session: CodingSessionAssessmentSummary instance (or any object with
                     a model_dump() method) containing the session context and
                     per-assessment feedback.

        Returns:
            Dict with ``message`` and ``id`` of the created report document.

        Raises:
            RuntimeError: If the proxy returns an error.
        """
        data = session.model_dump() if hasattr(session, "model_dump") else dict(session)
        return self._request("POST", "/reports/submit", data)

    def request_vertex_token(self) -> Dict[str, Any]:
        """Request a short-lived Vertex AI access token from the proxy.

        The proxy mints the token by impersonating a dedicated evaluator SA.
        Only available for EVALUATOR / ANCHOR roles when the LLM provider
        is Vertex-based.

        Returns:
            Dict with ``access_token`` (str) and ``expires_at`` (int, epoch seconds).

        Raises:
            RuntimeError: If the proxy refuses the request (wrong role, wrong
                provider, or server-side error).
        """
        return self._request("POST", "/config/vertex-token")
