from __future__ import annotations

import logging
import os
from typing import Any, Dict, List, Optional

import httpx
import requests

from .api import (
    DEFAULT_TIMEOUT,
    ahandle_response,
    apost_with_retries,
    handle_response,
    post_with_retries,
)
from .models import AnalyzeBatchItemResult, AnalyzeTextResult
from .errors import ItsAIError, ValidationError


logger = logging.getLogger("its_ai")


def _normalize_base_url(base_url: str) -> str:
    return base_url.rstrip("/")


class ItsAIClient:
    """Synchronous ITS-AI client.

    Parameters:
        api_key: API key string. If omitted, reads from env `ITS_AI_API_KEY`.
        base_url: Base URL of the API (`https://api.its-ai.org` by default).
        timeout: Default request timeout in seconds (overridable per-call).
        max_retries: Number of retries on 429/5xx (default: 3).

    Usage:
        >>> with ItsAIClient(api_key="sk_live_...") as client:
        ...     res = client.analyze_text("Hello world")
    """

    def __init__(
        self,
        api_key: Optional[str] = None,
        *,
        base_url: str = "https://api.its-ai.org",
        timeout: float = DEFAULT_TIMEOUT,
        max_retries: int = 3,
        session: Optional[requests.Session] = None,
    ) -> None:
        if api_key is None:
            self.api_key = os.getenv("ITS_AI_API_KEY", "")
        else:
            self.api_key = api_key
        if not isinstance(self.api_key, str) or not self.api_key:
            raise ValidationError(status_code=None, type="validation", message="api_key is required")

        self.base_url: str = _normalize_base_url(base_url)
        self.timeout: float = timeout
        self.max_retries: int = max(0, int(max_retries))
        self._session: requests.Session = session or requests.Session()

    # Context manager support
    def __enter__(self) -> "ItsAIClient":
        return self

    def __exit__(self, exc_type, exc, tb) -> None:  # type: ignore[override]
        self.close()

    def close(self) -> None:
        self._session.close()

    # Public API
    def analyze_text(
        self,
        text: str,
        deep_scan: bool = False,
        *,
        timeout: Optional[float] = None,
    ) -> AnalyzeTextResult:
        """Analyze a single text.

        Args:
            text: Input text (UTF-8 English).
            deep_scan: If True, the response may include `segmentation_tokens`.
            timeout: Optional per-call timeout.

        Returns:
            AnalyzeTextResult with answer and optional segmentation tokens.
        """
        if not isinstance(text, str) or not text.strip():
            raise ValidationError(status_code=None, type="validation", message="text must be a non-empty string")

        url = f"{self.base_url}/api/text"
        payload: Dict[str, Any] = {"text": text, "api_key": self.api_key, "deep_scan": bool(deep_scan)}
        response = post_with_retries(self._session, url, payload, timeout or self.timeout, self.max_retries)
        data = handle_response(response)

        if not isinstance(data, dict) or "answer" not in data:
            raise ItsAIError(status_code=response.status_code, type=None, message="Unexpected response shape", payload=data)

        answer = float(data["answer"])  # type: ignore[arg-type]
        segmentation = data.get("segmentation_tokens")
        if segmentation is not None:
            segmentation_list = [float(x) for x in segmentation]
        else:
            segmentation_list = None
        return AnalyzeTextResult(answer=answer, segmentation_tokens=segmentation_list)

    def analyze_batch(
        self,
        texts: List[str],
        deep_scan: bool = False,
        *,
        timeout: Optional[float] = None,
        max_batch_size: Optional[int] = None,
    ) -> List[AnalyzeBatchItemResult]:
        """Analyze multiple texts in one call.

        Args:
            texts: List of strings. Must be non-empty.
            deep_scan: If True, per-item segmentation tokens may appear.
            timeout: Optional per-call timeout.
            max_batch_size: Optional chunking; if provided, requests are chunked and results concatenated.

        Returns:
            List of AnalyzeBatchItemResult preserving input order.
        """
        if not isinstance(texts, list) or not texts:
            raise ValidationError(status_code=None, type="validation", message="texts must be a non-empty list of strings")
        for t in texts:
            if not isinstance(t, str) or not t:
                raise ValidationError(status_code=None, type="validation", message="each item in texts must be a non-empty string")

        if max_batch_size and max_batch_size > 0 and len(texts) > max_batch_size:
            results: List[AnalyzeBatchItemResult] = []
            for i in range(0, len(texts), max_batch_size):
                chunk = texts[i : i + max_batch_size]
                results.extend(self.analyze_batch(chunk, deep_scan=deep_scan, timeout=timeout))
            return results

        url = f"{self.base_url}/api/batch/"
        payload: Dict[str, Any] = {"texts": texts, "api_key": self.api_key, "deep_scan": bool(deep_scan)}
        response = post_with_retries(self._session, url, payload, timeout or self.timeout, self.max_retries)
        data = handle_response(response)

        if not isinstance(data, dict) or "answer" not in data or not isinstance(data["answer"], list):
            raise ItsAIError(status_code=response.status_code, type=None, message="Unexpected response shape", payload=data)

        items: List[AnalyzeBatchItemResult] = []
        for item in data["answer"]:
            if not isinstance(item, dict):
                raise ItsAIError(status_code=response.status_code, type=None, message="Unexpected item shape", payload=item)
            text = str(item.get("text", ""))
            deep_flag = bool(item.get("deep_scan", False))
            error = item.get("error")
            if error is not None and not isinstance(error, str):
                error = str(error)
            answer_val = item.get("answer")
            answer_f: Optional[float] = float(answer_val) if answer_val is not None else None
            segmentation = item.get("segmentation_tokens")
            seg_list = [float(x) for x in segmentation] if isinstance(segmentation, list) else None
            items.append(
                AnalyzeBatchItemResult(
                    text=text,
                    answer=answer_f,
                    segmentation_tokens=seg_list,
                    deep_scan=deep_flag,
                    error=error,
                )
            )
        return items


class ItsAIAsyncClient:
    """Asynchronous ITS-AI client using httpx.AsyncClient."""

    def __init__(
        self,
        api_key: Optional[str] = None,
        *,
        base_url: str = "https://api.its-ai.org",
        timeout: float = DEFAULT_TIMEOUT,
        max_retries: int = 3,
        client: Optional[httpx.AsyncClient] = None,
    ) -> None:
        if api_key is None:
            self.api_key = os.getenv("ITS_AI_API_KEY", "")
        else:
            self.api_key = api_key
        if not isinstance(self.api_key, str) or not self.api_key:
            raise ValidationError(status_code=None, type="validation", message="api_key is required")

        self.base_url: str = _normalize_base_url(base_url)
        self.timeout: float = timeout
        self.max_retries: int = max(0, int(max_retries))
        self._client: httpx.AsyncClient = client or httpx.AsyncClient()

    async def __aenter__(self) -> "ItsAIAsyncClient":  # type: ignore[override]
        return self

    async def __aexit__(self, exc_type, exc, tb) -> None:  # type: ignore[override]
        await self.aclose()

    async def aclose(self) -> None:
        await self._client.aclose()

    async def analyze_text(
        self,
        text: str,
        deep_scan: bool = False,
        *,
        timeout: Optional[float] = None,
    ) -> AnalyzeTextResult:
        if not isinstance(text, str) or not text.strip():
            raise ValidationError(status_code=None, type="validation", message="text must be a non-empty string")

        url = f"{self.base_url}/api/text"
        payload: Dict[str, Any] = {"text": text, "api_key": self.api_key, "deep_scan": bool(deep_scan)}
        response = await apost_with_retries(self._client, url, payload, timeout or self.timeout, self.max_retries)
        data = await ahandle_response(response)

        if not isinstance(data, dict) or "answer" not in data:
            raise ItsAIError(status_code=response.status_code, type=None, message="Unexpected response shape", payload=data)  # type: ignore[arg-type]

        answer = float(data["answer"])  # type: ignore[arg-type]
        segmentation = data.get("segmentation_tokens")
        if segmentation is not None:
            segmentation_list = [float(x) for x in segmentation]
        else:
            segmentation_list = None
        return AnalyzeTextResult(answer=answer, segmentation_tokens=segmentation_list)

    async def analyze_batch(
        self,
        texts: List[str],
        deep_scan: bool = False,
        *,
        timeout: Optional[float] = None,
        max_batch_size: Optional[int] = None,
    ) -> List[AnalyzeBatchItemResult]:
        if not isinstance(texts, list) or not texts:
            raise ValidationError(status_code=None, type="validation", message="texts must be a non-empty list of strings")
        for t in texts:
            if not isinstance(t, str) or not t:
                raise ValidationError(status_code=None, type="validation", message="each item in texts must be a non-empty string")

        if max_batch_size and max_batch_size > 0 and len(texts) > max_batch_size:
            results: List[AnalyzeBatchItemResult] = []
            for i in range(0, len(texts), max_batch_size):
                chunk = texts[i : i + max_batch_size]
                results.extend(
                    await self.analyze_batch(chunk, deep_scan=deep_scan, timeout=timeout)
                )
            return results

        url = f"{self.base_url}/api/batch/"
        payload: Dict[str, Any] = {"texts": texts, "api_key": self.api_key, "deep_scan": bool(deep_scan)}
        response = await apost_with_retries(self._client, url, payload, timeout or self.timeout, self.max_retries)
        data = await ahandle_response(response)

        if not isinstance(data, dict) or "answer" not in data or not isinstance(data["answer"], list):
            raise ItsAIError(status_code=response.status_code, type=None, message="Unexpected response shape", payload=data)  # type: ignore[arg-type]

        items: List[AnalyzeBatchItemResult] = []
        for item in data["answer"]:
            if not isinstance(item, dict):
                raise ItsAIError(status_code=response.status_code, type=None, message="Unexpected item shape", payload=item)  # type: ignore[arg-type]
            text = str(item.get("text", ""))
            deep_flag = bool(item.get("deep_scan", False))
            error = item.get("error")
            if error is not None and not isinstance(error, str):
                error = str(error)
            answer_val = item.get("answer")
            answer_f: Optional[float] = float(answer_val) if answer_val is not None else None
            segmentation = item.get("segmentation_tokens")
            seg_list = [float(x) for x in segmentation] if isinstance(segmentation, list) else None
            items.append(
                AnalyzeBatchItemResult(
                    text=text,
                    answer=answer_f,
                    segmentation_tokens=seg_list,
                    deep_scan=deep_flag,
                    error=error,
                )
            )
        return items


