"""
ZeliAI Python SDK — 16개+ AI 모델을 하나의 API로

Usage:
    from zeliai import ZeliAI

    client = ZeliAI(api_key="zeliai-xxxx...")
    response = client.chat("안녕하세요!", model="gpt-5.2")
    print(response.content)
"""

import os
import json
from typing import Optional, List, Dict, Generator, Any

import httpx

from .models import (
    ChatResponse, ChatMessage, ChatChoice, UsageInfo,
    SearchResponse, SearchResult,
    DocQAResponse, DocSource,
    ModelInfo, AccountUsage,
)
from .exceptions import (
    ZeliAIError, AuthenticationError, RateLimitError,
    ModelNotFoundError, InvalidRequestError, ServerError,
)

__all__ = ["ZeliAI"]

# 기본 B2B API URL
DEFAULT_BASE_URL = "https://api.zeliai.com"
API_PREFIX = "/api/b2b/v1"


class ZeliAI:
    """
    ZeliAI B2B API 클라이언트

    Args:
        api_key: ZeliAI API 키 (zeliai-xxxxx 형식).
                 환경변수 ZELIAI_API_KEY 로도 설정 가능.
        base_url: API 서버 URL. 기본값: https://api.zeliai.com
        timeout: 요청 타임아웃 (초). 기본값: 120
    """

    def __init__(
        self,
        api_key: Optional[str] = None,
        base_url: Optional[str] = None,
        timeout: float = 120.0,
    ):
        self.api_key = api_key or os.environ.get("ZELIAI_API_KEY")
        if not self.api_key:
            raise AuthenticationError(
                "API key is required. Pass api_key='zeliai-...' "
                "or set ZELIAI_API_KEY environment variable."
            )

        self.base_url = (base_url or os.environ.get("ZELIAI_BASE_URL", DEFAULT_BASE_URL)).rstrip("/")
        self._client = httpx.Client(
            base_url=self.base_url + API_PREFIX,
            headers={
                "Authorization": f"Bearer {self.api_key}",
                "Content-Type": "application/json",
                "User-Agent": "zeliai-python/0.1.0",
            },
            timeout=timeout,
        )

    def __del__(self):
        try:
            self._client.close()
        except Exception:
            pass

    def __enter__(self):
        return self

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

    # ═══════════════════════════════════════════════════
    #  Chat
    # ═══════════════════════════════════════════════════

    def chat(
        self,
        message: str,
        model: str = "gpt-5.2",
        system: Optional[str] = None,
        session_id: Optional[str] = None,
        temperature: Optional[float] = None,
        max_tokens: Optional[int] = None,
    ) -> ChatResponse:
        """
        채팅 완성 (동기 응답)

        Args:
            message: 사용자 메시지
            model: 사용할 모델 ID (예: gpt-5.2, claude-opus-4.5, gemini-3-pro)
            system: 시스템 지시 (선택)
            session_id: 세션 ID (대화 이어가기)
            temperature: 생성 온도 (0.0 ~ 2.0)
            max_tokens: 최대 생성 토큰 수

        Returns:
            ChatResponse: 응답 객체 (.content 로 텍스트 접근)

        Example:
            >>> response = client.chat("파이썬으로 정렬 알고리즘 설명해줘")
            >>> print(response.content)
        """
        messages = self._build_messages(message, system)
        body = {
            "model": model,
            "messages": messages,
            "stream": False,
        }
        if session_id:
            body["session_id"] = session_id
        if temperature is not None:
            body["temperature"] = temperature
        if max_tokens is not None:
            body["max_tokens"] = max_tokens

        data = self._post("/chat/completions", body)
        return self._parse_chat_response(data)

    def chat_stream(
        self,
        message: str,
        model: str = "gpt-5.2",
        system: Optional[str] = None,
        session_id: Optional[str] = None,
        temperature: Optional[float] = None,
        max_tokens: Optional[int] = None,
    ) -> Generator[str, None, None]:
        """
        채팅 스트리밍 (SSE)

        실시간으로 응답 텍스트를 한 조각씩 yield 합니다.

        Args:
            message: 사용자 메시지
            model: 사용할 모델 ID
            system: 시스템 지시 (선택)
            session_id: 세션 ID
            temperature: 생성 온도
            max_tokens: 최대 생성 토큰 수

        Yields:
            str: 응답 텍스트 조각

        Example:
            >>> for chunk in client.chat_stream("스트리밍 테스트"):
            ...     print(chunk, end="", flush=True)
        """
        messages = self._build_messages(message, system)
        body = {
            "model": model,
            "messages": messages,
            "stream": True,
        }
        if session_id:
            body["session_id"] = session_id
        if temperature is not None:
            body["temperature"] = temperature
        if max_tokens is not None:
            body["max_tokens"] = max_tokens

        with self._client.stream(
            "POST",
            "/chat/completions",
            json=body,
        ) as response:
            self._check_response_error(response)
            for line in response.iter_lines():
                if not line.startswith("data: "):
                    continue
                payload = line[6:]
                if payload.strip() == "[DONE]":
                    return
                try:
                    chunk = json.loads(payload)
                    delta = chunk.get("choices", [{}])[0].get("delta", {})
                    content = delta.get("content")
                    if content:
                        yield content
                except (json.JSONDecodeError, IndexError, KeyError):
                    continue

    # ═══════════════════════════════════════════════════
    #  Search
    # ═══════════════════════════════════════════════════

    def search(
        self,
        query: str,
        max_results: int = 5,
    ) -> SearchResponse:
        """
        AI 기반 웹 검색

        Args:
            query: 검색 쿼리
            max_results: 최대 결과 수 (기본 5)

        Returns:
            SearchResponse: .results (리스트) + .ai_summary (AI 요약)

        Example:
            >>> result = client.search("최신 AI 트렌드 2026")
            >>> print(result.ai_summary)
            >>> for r in result.results:
            ...     print(f"  - {r.title}: {r.url}")
        """
        body = {"query": query, "max_results": max_results}
        data = self._post("/search", body)
        return self._parse_search_response(data)

    # ═══════════════════════════════════════════════════
    #  Document QA
    # ═══════════════════════════════════════════════════

    def document_qa(
        self,
        question: str,
        file_path: str,
        model: str = "gpt-5.2",
    ) -> DocQAResponse:
        """
        문서 기반 질의응답

        파일을 업로드하고 해당 문서에 대한 질문을 합니다.

        Args:
            question: 문서에 대한 질문
            file_path: 업로드할 파일 경로
            model: 사용할 모델 ID

        Returns:
            DocQAResponse: .answer (답변 텍스트)

        Example:
            >>> result = client.document_qa("이 계약서의 만료일은?", "contract.pdf")
            >>> print(result.answer)
        """
        with open(file_path, "rb") as f:
            files = {"file": (os.path.basename(file_path), f)}
            form_data = {"question": question, "model": model}

            response = self._client.post(
                "/documents/qa",
                data=form_data,
                files=files,
                headers={"Content-Type": None},  # let httpx set multipart
            )
        self._check_response_error(response)
        data = response.json()
        return self._parse_docqa_response(data)

    # ═══════════════════════════════════════════════════
    #  Utility endpoints
    # ═══════════════════════════════════════════════════

    def models(self) -> List[ModelInfo]:
        """
        사용 가능한 모델 목록 조회

        Returns:
            List[ModelInfo]: 모델 정보 리스트

        Example:
            >>> for m in client.models():
            ...     print(f"{m.id}: {m.description}")
        """
        data = self._get("/models")
        result = []
        for item in data.get("data", []):
            result.append(ModelInfo(
                id=item.get("id", ""),
                name=item.get("name", ""),
                provider=item.get("provider", ""),
                description=item.get("description", ""),
                max_tokens=item.get("max_tokens", 0),
                supports_streaming=item.get("supports_streaming", True),
                supports_search=item.get("supports_search", False),
            ))
        return result

    def usage(self) -> AccountUsage:
        """
        현재 API 키 사용량 조회

        Returns:
            AccountUsage: 사용량 정보

        Example:
            >>> u = client.usage()
            >>> print(f"총 요청: {u.total_requests}, 총 토큰: {u.total_tokens}")
        """
        data = self._get("/usage")
        return AccountUsage(
            company_name=data.get("company_name", ""),
            plan=data.get("plan", ""),
            total_requests=data.get("total_requests", 0),
            total_tokens=data.get("total_tokens", 0),
            rpm_limit=data.get("rpm_limit", 0),
            tpm_limit=data.get("tpm_limit", 0),
            is_active=data.get("is_active", True),
        )

    def health(self) -> Dict[str, Any]:
        """
        서버 상태 확인

        Returns:
            dict: 서버 상태 정보

        Example:
            >>> status = client.health()
            >>> print(status)  # {"status": "ok", ...}
        """
        return self._get("/health")

    # ═══════════════════════════════════════════════════
    #  Internal helpers
    # ═══════════════════════════════════════════════════

    @staticmethod
    def _build_messages(message: str, system: Optional[str] = None) -> List[Dict[str, str]]:
        msgs = []
        if system:
            msgs.append({"role": "system", "content": system})
        msgs.append({"role": "user", "content": message})
        return msgs

    def _get(self, path: str) -> dict:
        response = self._client.get(path)
        self._check_response_error(response)
        return response.json()

    def _post(self, path: str, body: dict) -> dict:
        response = self._client.post(path, json=body)
        self._check_response_error(response)
        return response.json()

    @staticmethod
    def _check_response_error(response):
        if response.status_code == 200 or response.status_code == 201:
            return

        try:
            data = response.json()
            error = data.get("error", {})
            msg = error.get("message", response.text)
            err_type = error.get("type", "api_error")
            err_code = error.get("code")
        except Exception:
            msg = response.text
            err_type = "api_error"
            err_code = None

        status = response.status_code
        if status == 401:
            raise AuthenticationError(msg)
        elif status == 429:
            raise RateLimitError(msg)
        elif status == 404 and err_code == "model_not_found":
            raise ModelNotFoundError(msg)
        elif status == 400:
            raise InvalidRequestError(msg)
        elif status >= 500:
            raise ServerError(msg)
        else:
            raise ZeliAIError(msg, status_code=status, error_type=err_type, error_code=err_code)

    @staticmethod
    def _parse_chat_response(data: dict) -> ChatResponse:
        choices = []
        for c in data.get("choices", []):
            msg = c.get("message", {})
            choices.append(ChatChoice(
                index=c.get("index", 0),
                message=ChatMessage(
                    role=msg.get("role", "assistant"),
                    content=msg.get("content", ""),
                ),
                finish_reason=c.get("finish_reason", "stop"),
            ))

        usage_data = data.get("usage", {})
        return ChatResponse(
            id=data.get("id", ""),
            object=data.get("object", "chat.completion"),
            created=data.get("created", 0),
            model=data.get("model", ""),
            choices=choices,
            usage=UsageInfo(
                prompt_tokens=usage_data.get("prompt_tokens", 0),
                completion_tokens=usage_data.get("completion_tokens", 0),
                total_tokens=usage_data.get("total_tokens", 0),
            ),
            session_id=data.get("session_id"),
        )

    @staticmethod
    def _parse_search_response(data: dict) -> SearchResponse:
        results = []
        for r in data.get("results", []):
            results.append(SearchResult(
                title=r.get("title", ""),
                url=r.get("url", ""),
                snippet=r.get("snippet", ""),
                source=r.get("source", "web"),
            ))
        return SearchResponse(
            id=data.get("id", ""),
            query=data.get("query", ""),
            results=results,
            ai_summary=data.get("ai_summary"),
            total_results=data.get("total_results", 0),
        )

    @staticmethod
    def _parse_docqa_response(data: dict) -> DocQAResponse:
        sources = []
        for s in data.get("sources", []):
            sources.append(DocSource(
                file_name=s.get("file_name", ""),
                relevance=s.get("relevance", 0.0),
            ))
        return DocQAResponse(
            id=data.get("id", ""),
            answer=data.get("answer", ""),
            model=data.get("model", ""),
            sources=sources,
        )
