# client.py — Sync and async clients for the COS API.
#
# Usage (sync):
#   from cos_sdk import COS
#   client = COS(api_key="cos_live_xxx")
#   result = client.validate("Some AI-generated text", tier="tier_2")
#   print(result.confidence_score, result.risk_level)
#
# Usage (async):
#   from cos_sdk import AsyncCOS
#   client = AsyncCOS(api_key="cos_live_xxx")
#   result = await client.validate("Some AI-generated text")
#
# Usage (streaming):
#   for event in client.validate_stream("Some text", tier="tier_3"):
#       if event.event == "t1_result":
#           print("Quick check:", event.result.confidence_score)
#       elif event.event == "validation_complete":
#           print("Full result:", event.result.confidence_score)

from typing import AsyncIterator, Iterator, Optional

import httpx

from cos_sdk.models import ValidationResult, StreamEvent
from cos_sdk.streaming import parse_sse_stream
from cos_sdk.exceptions import (
    COSError, AuthenticationError, ForbiddenError,
    RateLimitError, ValidationError, ServerError,
)

# Default COS API base URL
DEFAULT_BASE_URL = "https://cos.protofine.ai"


def _handle_error(response: httpx.Response):
    """Check for HTTP errors and raise the appropriate SDK exception."""
    if response.status_code == 200:
        return

    try:
        body = response.json()
    except Exception:
        body = {"message": response.text}

    # Extract error message from the response
    if isinstance(body, dict) and "detail" in body:
        detail = body["detail"]
        if isinstance(detail, dict):
            message = detail.get("message", str(detail))
        else:
            message = str(detail)
    else:
        message = str(body)

    if response.status_code == 401:
        raise AuthenticationError(message, status_code=401, response_body=body)
    elif response.status_code == 403:
        raise ForbiddenError(message, status_code=403, response_body=body)
    elif response.status_code == 429:
        retry_after = 60
        if isinstance(body, dict) and "detail" in body and isinstance(body["detail"], dict):
            retry_after = body["detail"].get("retry_after", 60)
        raise RateLimitError(message, retry_after=retry_after, status_code=429, response_body=body)
    elif response.status_code == 400 or response.status_code == 422:
        raise ValidationError(message, status_code=response.status_code, response_body=body)
    elif response.status_code >= 500:
        raise ServerError(message, status_code=response.status_code, response_body=body)
    else:
        raise COSError(message, status_code=response.status_code, response_body=body)


class COS:
    """Synchronous client for the COS API.

    Example:
        client = COS(api_key="cos_live_xxx")
        result = client.validate("AI-generated text to check")
        if result.risk_level == "high":
            print("Warning:", result.flagged_claims)
    """

    def __init__(self, api_key: str, base_url: str = DEFAULT_BASE_URL, timeout: float = 30.0):
        """Initialize the COS client.

        Args:
            api_key: Your COS API key (starts with cos_live_)
            base_url: COS API base URL (default: https://cos.protofine.ai)
            timeout: Request timeout in seconds (default: 30)
        """
        if not api_key or not api_key.startswith("cos_live_"):
            raise ValueError("API key must start with 'cos_live_'. Get one at https://cos.protofine.ai/console")

        self.api_key = api_key
        self.base_url = base_url.rstrip("/")
        self._client = httpx.Client(
            base_url=self.base_url,
            headers={
                "Authorization": f"Bearer {api_key}",
                "Content-Type": "application/json",
            },
            timeout=timeout,
        )

    def validate(self, text: str, tier: str = "tier_2", context: Optional[str] = None) -> ValidationResult:
        """Validate AI-generated text for hallucinations.

        Args:
            text: The AI-generated text to check
            tier: Validation tier — "tier_1" (instant), "tier_2" (AI review), "tier_3" (multi-model)
            context: Original prompt or context (optional, improves accuracy)

        Returns:
            ValidationResult with confidence_score, risk_level, and flagged_claims
        """
        body = {"text": text, "tier": tier}
        if context:
            body["context"] = context

        response = self._client.post("/api/v2/validate", json=body)
        _handle_error(response)
        return ValidationResult.from_dict(response.json())

    def validate_stream(self, text: str, tier: str = "tier_2",
                        context: Optional[str] = None) -> Iterator[StreamEvent]:
        """Validate AI text with streaming results (SSE).

        Yields events progressively:
        - t1_result: Instant heuristic check (~5ms)
        - validation_complete: Full result with corrections
        - done: Stream complete

        Args:
            text: The AI-generated text to check
            tier: Validation tier
            context: Original prompt or context (optional)

        Yields:
            StreamEvent objects with event type and data
        """
        body = {"text": text, "tier": tier}
        if context:
            body["context"] = context

        with self._client.stream("POST", "/api/v2/validate/stream", json=body) as response:
            _handle_error(response)
            yield from parse_sse_stream(response.iter_lines())

    def close(self):
        """Close the underlying HTTP connection."""
        self._client.close()

    def __enter__(self):
        return self

    def __exit__(self, *args):
        self.close()


class AsyncCOS:
    """Asynchronous client for the COS API.

    Example:
        async with AsyncCOS(api_key="cos_live_xxx") as client:
            result = await client.validate("AI-generated text")
            print(result.confidence_score)
    """

    def __init__(self, api_key: str, base_url: str = DEFAULT_BASE_URL, timeout: float = 30.0):
        if not api_key or not api_key.startswith("cos_live_"):
            raise ValueError("API key must start with 'cos_live_'. Get one at https://cos.protofine.ai/console")

        self.api_key = api_key
        self.base_url = base_url.rstrip("/")
        self._client = httpx.AsyncClient(
            base_url=self.base_url,
            headers={
                "Authorization": f"Bearer {api_key}",
                "Content-Type": "application/json",
            },
            timeout=timeout,
        )

    async def validate(self, text: str, tier: str = "tier_2",
                       context: Optional[str] = None) -> ValidationResult:
        """Validate AI-generated text (async version)."""
        body = {"text": text, "tier": tier}
        if context:
            body["context"] = context

        response = await self._client.post("/api/v2/validate", json=body)
        _handle_error(response)
        return ValidationResult.from_dict(response.json())

    async def validate_stream(self, text: str, tier: str = "tier_2",
                              context: Optional[str] = None) -> AsyncIterator[StreamEvent]:
        """Validate AI text with streaming results (async SSE).

        Yields events progressively:
        - t1_result: Instant heuristic check (~5ms)
        - validation_complete: Full result with corrections
        - done: Stream complete

        Args:
            text: The AI-generated text to check
            tier: Validation tier
            context: Original prompt or context (optional)

        Yields:
            StreamEvent objects with event type and data
        """
        body = {"text": text, "tier": tier}
        if context:
            body["context"] = context

        async with self._client.stream("POST", "/api/v2/validate/stream", json=body) as response:
            _handle_error(response)
            event_type = ""
            data_lines = []
            async for line in response.aiter_lines():
                line = line.rstrip("\n").rstrip("\r")
                if line.startswith("event: "):
                    event_type = line[7:]
                elif line.startswith("data: "):
                    data_lines.append(line[6:])
                elif line == "":
                    if event_type and data_lines:
                        import json
                        raw_data = "\n".join(data_lines)
                        try:
                            data = json.loads(raw_data)
                        except json.JSONDecodeError:
                            data = {"raw": raw_data}
                        yield StreamEvent(event=event_type, data=data)
                        if event_type == "done":
                            return
                    event_type = ""
                    data_lines = []

    async def close(self):
        """Close the underlying HTTP connection."""
        await self._client.aclose()

    async def __aenter__(self):
        return self

    async def __aexit__(self, *args):
        await self.close()
