from __future__ import annotations

import os
from dataclasses import dataclass
from typing import Any

import requests

from .exceptions import (
    AuthenticationError,
    BaudLabsError,
    InsufficientCreditsError,
    NotFoundError,
    ServerError,
    ValidationError,
)

DEFAULT_BASE_URL = "https://api.baudlabs.ai"
VALID_MEDIA_TYPES = {"image", "video", "mixed"}


@dataclass(slots=True)
class Index:
    id: str
    status: str = "ready"
    asset_count: int = 0


@dataclass(slots=True)
class Match:
    score: float
    image_id: str | None = None
    video_id: str | None = None
    start_time: float | None = None
    end_time: float | None = None
    image_url: str | None = None
    video_url: str | None = None


@dataclass(slots=True)
class SearchResult:
    matches: list[Match]
    query: str | None = None
    credits_remaining: float | None = None

    def __getitem__(self, key: str):
        if not hasattr(self, key):
            raise KeyError(key)
        return getattr(self, key)

    def get(self, key: str, default=None):
        return getattr(self, key, default)

    def keys(self):
        return ("matches", "query", "credits_remaining")

    def values(self):
        return tuple(getattr(self, key) for key in self.keys())

    def items(self):
        return tuple((key, getattr(self, key)) for key in self.keys())

    def __contains__(self, key: str) -> bool:
        return key in self.keys()


class _HTTPClient:
    """Thin wrapper around ``requests`` that injects auth and maps errors."""

    def __init__(self, api_key: str, base_url: str, timeout: float):
        self.base_url = base_url.rstrip("/")
        self.timeout = timeout
        self.session = requests.Session()
        self.session.headers.update({"X-API-Key": api_key})

    @staticmethod
    def _json_body(resp: requests.Response) -> dict:
        try:
            return resp.json()
        except ValueError:
            return {}

    def _handle_response(self, resp: requests.Response) -> dict:
        body = self._json_body(resp)
        if resp.status_code == 401:
            raise AuthenticationError(
                body.get("detail", "Invalid API key"),
                status_code=401,
                response=body,
            )
        if resp.status_code == 402:
            raise InsufficientCreditsError(
                body.get("error", "Insufficient credits"),
                credits_remaining=body.get("credits_remaining", 0),
                cost=body.get("cost", 0),
                status_code=402,
                response=body,
            )
        if resp.status_code == 404:
            raise NotFoundError(
                body.get("detail", "Not found"),
                status_code=404,
                response=body,
            )
        if resp.status_code in {400, 422}:
            raise ValidationError(
                str(body.get("detail", "Validation error")),
                status_code=resp.status_code,
                response=body,
            )
        if resp.status_code >= 500:
            raise ServerError(f"Server error ({resp.status_code})", status_code=resp.status_code)
        if not resp.ok:
            raise BaudLabsError(f"Request failed ({resp.status_code})", status_code=resp.status_code)
        return body

    def get(self, path: str, **kwargs: Any) -> dict:
        resp = self.session.get(f"{self.base_url}{path}", timeout=self.timeout, **kwargs)
        return self._handle_response(resp)

    def post(self, path: str, **kwargs: Any) -> dict:
        resp = self.session.post(f"{self.base_url}{path}", timeout=self.timeout, **kwargs)
        return self._handle_response(resp)


def _resolve_file(path_or_bytes, name: str | None = None):
    """Accept a file path (str/Path) or raw bytes and return (basename, file-like)."""
    if isinstance(path_or_bytes, (str, os.PathLike)):
        path_or_bytes = str(path_or_bytes)
        return os.path.basename(path_or_bytes), open(path_or_bytes, "rb")
    if isinstance(path_or_bytes, bytes):
        import io

        return name or "file", io.BytesIO(path_or_bytes)
    return name or "file", path_or_bytes


def _normalize_media_type(media_type: str) -> str:
    value = (media_type or "").strip().lower()
    if value not in VALID_MEDIA_TYPES:
        allowed = ", ".join(f'"{item}"' for item in sorted(VALID_MEDIA_TYPES))
        raise ValidationError(f"media_type must be one of {allowed}.")
    return value


class Client:
    """BaudLabs search client."""

    def __init__(
        self,
        api_key: str | None = None,
        *,
        base_url: str = DEFAULT_BASE_URL,
        timeout: float = 120.0,
    ):
        api_key = api_key or os.environ.get("BAUDLABS_API_KEY", "")
        if not api_key:
            raise AuthenticationError(
                "No API key provided. Pass it to Client() or set BAUDLABS_API_KEY."
            )
        self._http = _HTTPClient(api_key, base_url, timeout)

    def create_index(
        self,
        index_id: str,
        source: str | None = None,
    ) -> Index:
        """Create a search index."""
        if source is not None:
            raise ValidationError(
                "source is not supported by the current public API. "
                "Create the index first, then upload local files with add_images() or add_videos()."
            )
        resp = self._http.post("/v1/create-index", json={"index_id": index_id})
        index = resp.get("index", {})
        return Index(
            id=index.get("index_name", index_id),
            status="ready",
            asset_count=index.get("asset_count", 0),
        )

    def add_images(self, index_id: str, files: list[str]) -> dict:
        """Upload local image files into an existing search index."""
        return self._upload_files(
            path="/v1/add-images",
            index_id=index_id,
            files=files,
        )

    def add_videos(self, index_id: str, files: list[str]) -> dict:
        """Upload local video files into an existing search index."""
        return self._upload_files(
            path="/v1/add-videos",
            index_id=index_id,
            files=files,
        )

    def search(
        self,
        index_id: str,
        query: str | None = None,
        image_query: str | None = None,
        media_type: str = "mixed",
        top_k: int = 10,
    ) -> SearchResult:
        """Search an index with a text or image query."""
        media_type = _normalize_media_type(media_type)
        search_text = self._resolve_search_text(query=query, image_query=image_query)
        top_k = int(top_k)
        if top_k <= 0:
            raise ValidationError("top_k must be greater than 0.")
        resp = self._http.post(
            "/v1/search",
            json={
                "index_id": index_id,
                "query": query,
                "image_query": image_query,
                "media_type": media_type,
                "top_k": top_k,
            },
        )
        matches = [
            Match(
                image_id=item.get("image_id"),
                video_id=item.get("video_id"),
                image_url=item.get("image_url"),
                video_url=item.get("video_url"),
                score=float(item.get("score", 0.0)),
                start_time=item.get("start_time"),
                end_time=item.get("end_time"),
            )
            for item in resp.get("results", [])
        ]
        return SearchResult(
            matches=matches,
            query=resp.get("query", search_text),
            credits_remaining=resp.get("credits_remaining"),
        )

    def _upload_files(
        self,
        *,
        path: str,
        index_id: str,
        files: list[str],
        extra_data: dict[str, str] | None = None,
    ) -> dict:
        if not files:
            raise ValidationError("files must contain at least one path.")

        upload_files = []
        opened = []
        try:
            for item in files:
                fname, fobj = _resolve_file(item)
                opened.append(fobj)
                upload_files.append(("files", (fname, fobj)))

            data = {"index_id": index_id}
            if extra_data:
                data.update(extra_data)
            return self._http.post(path, files=upload_files, data=data)
        finally:
            for fobj in opened:
                if hasattr(fobj, "close"):
                    fobj.close()

    @staticmethod
    def _resolve_search_text(query: str | None, image_query: str | None) -> str:
        if bool(query) == bool(image_query):
            raise ValidationError("Provide exactly one of query or image_query.")
        return (query or image_query or "").strip()
