# client.py
import os
import httpx
import zipfile
import io
import json
import time
from pathlib import Path
from typing import Any

from poma.exceptions import ApiError, InvalidInputError, InvalidResponseError, PomaSDKError
from poma.retrieval import generate_cheatsheets


USER_AGENT = "poma-ai-sdk/0.1.0"

API_BASE_URL = "https://api.poma-ai.com/api/v1"


def extract_chunks_and_chunksets_from_poma_archive(
    poma_archive_data: bytes | None = None,
    poma_archive_path: str | os.PathLike[str] | None = None,
) -> dict[str, Any]:
    """
    Extract chunks and chunksets from a POMA archive file.
    POMA archive is a zip file containing chunks.json and chunksets.json.
    Args:
        poma_archive_data: The POMA archive as bytes.
        poma_archive_path: Path to the POMA archive file.
    Returns:
        dict: A dictionary with ``chunks`` and ``chunksets`` keys.
    """
    chunks = None
    chunksets = None
    if poma_archive_path:
        with zipfile.ZipFile(poma_archive_path, "r") as zip_ref:
            chunks = zip_ref.read("chunks.json")
            chunksets = zip_ref.read("chunksets.json")
    elif poma_archive_data:
        with zipfile.ZipFile(io.BytesIO(poma_archive_data), "r") as zip_ref:
            chunks = zip_ref.read("chunks.json")
            chunksets = zip_ref.read("chunksets.json")
    else:
        raise ValueError(
            "Either poma_archive_data or poma_archive_path must be provided."
        )
    if not chunks or not chunksets:
        raise KeyError("Result must contain 'chunks' and 'chunksets' keys.")
    return {"chunks": json.loads(chunks), "chunksets": json.loads(chunksets)}


class Poma:
    """
    Client for interacting with the POMA API.
    """

    def __init__(
        self,
        api_key: str | None = None,
        *,
        timeout: float = 600.0,
        client: httpx.Client | None = None,
    ):
        """
        Initialize the POMA client.
        Args:
            api_key (str, optional):
                API key for authenticating with POMA. If not provided,
                the value is read from the environment variable `POMA_API_KEY`.
            timeout (float, default=600.0):
                Timeout (in seconds) for all HTTP requests.
            client (httpx.Client, optional):
                Custom HTTP client. If not provided, a default client is created.
        """
        api_base_url = API_BASE_URL
        # Override API base URL if environment variable is set
        if os.environ.get("API_BASE_URL"):
            api_base_url = os.environ.get("API_BASE_URL")
        if not api_base_url:
            raise ValueError("API base URL cannot be empty.")

        self.base_api_url = api_base_url.rstrip("/")
        self._client = client or httpx.Client(
            timeout=timeout, headers={"user-agent": USER_AGENT}
        )
        if not (api_key := api_key or os.environ.get("POMA_API_KEY", "")):
            raise Exception("POMA_API_KEY environment variable not set.")
        self._client.headers.update({"Authorization": f"Bearer {api_key}"})

    def start_chunk_file(
        self,
        file_path: str | os.PathLike[str],
        *,
        base_url: str | None = None,
    ) -> dict[str, Any]:
        """
        Submit a file with text to POMA for chunking.

        Uses ``POST /ingest``. On HTTP error, raises :exc:`ApiError` with
        :attr:`ApiError.status_code` and :attr:`ApiError.response_body` set from
        the API response. Typical status codes: 400 (bad request / failed to create job),
        401 (missing or invalid authorization), 403 (not found, quota exceeded, too many jobs),
        500 (internal error, failed to upload file).

        Args:
            file_path (str | os.PathLike[str]):
                Path to the input file (string or path-like). Must have an allowed file extension.
            base_url (str, optional):
                Optional base URL to resolve relative links within the file.
        Returns:
            A dictionary containing a unique job identifier for the submitted job.

        Raises:
            ApiError: On API HTTP error (see :attr:`ApiError.status_code` and
                :attr:`ApiError.response_body`).
            InvalidResponseError: If the response body is not valid JSON.
        """
        if file_path is None or (isinstance(file_path, str) and not file_path.strip()):
            raise ValueError("file_path must be a non-empty string or path-like.")
        path = Path(file_path)
        payload = {}
        payload["is_sdk"] = True
        if base_url:
            payload["base_url"] = base_url
        try:
            response = self._client.post(
                f"{self.base_api_url}/ingest",
                data=payload,
                files={
                    "file": (path.name, path.read_bytes()),
                },
            )
            response.raise_for_status()
        except httpx.HTTPStatusError as error:
            status = error.response.status_code
            body = (error.response.text or "").strip()
            msg = f"Failed to submit file '{path}'"
            if body:
                msg = f"{msg}: {body}"
            raise ApiError(
                msg, status_code=status, response_body=body, file_path=str(path)
            ) from error
        try:
            data = response.json()
        except ValueError as error:
            raise InvalidResponseError(
                f"Server returned non-JSON or empty body (file: {path})"
            ) from error
        return data

    def get_chunk_result(
        self,
        job_id: str,
        *,
        initial_delay: float = 5.0,
        poll_interval: float = 3.0,
        max_interval: float = 15.0,
        show_progress: bool = False,
        download_dir: str | os.PathLike[str] | None = None,
        filename: str | None = None,
    ) -> dict[str, Any]:
        """
        Poll POMA for the result of a chunking job until completion.

        Uses ``GET /jobs/{job_id}/status``. On HTTP error, raises :exc:`ApiError` with
        :attr:`ApiError.status_code` and :attr:`ApiError.response_body` set from the API.
        Typical status codes: 401 (authorization), 403 (account not found), 404 (job not found
        or no access), 500 (failed to get job result). If the job completes with status
        ``"failed"``, also raises :exc:`ApiError` (without status_code/response_body).

        Args:
            job_id (str):
                The unique identifier of the submitted job.
            initial_delay (float, default=5.0):
                Initial delay (in seconds) before the first poll request.
            poll_interval (float, default=1.0):
                Starting interval (in seconds) between polling requests.
            max_interval (float, default=15.0):
                Maximum interval (in seconds) between polling requests.
            show_progress (bool, default=False):
                If True, logs progress messages during polling.
            download_dir (str | os.PathLike[str], optional):
                Directory to save the downloaded file in. If neither download_dir nor
                filename is set, the result is returned in memory (no file saved). If
                filename is set but download_dir is not, the file is saved in the
                current directory.
            filename (str, optional):
                Name for the saved .poma file. If it does not end with ``.poma``, that
                suffix is appended. If not set when saving to disk, uses the server
                filename when provided, otherwise ``{job_id}.poma``.
        Returns:
            The JSON result containing at least the keys `chunks` and `chunksets`.

        Raises:
            ApiError: On API HTTP error or when job status is ``"failed"`` (see
                :attr:`ApiError.status_code` and :attr:`ApiError.response_body` where set).
            InvalidResponseError: On unexpected job status in the response.
        """
        time.sleep(initial_delay)
        current_interval = poll_interval
        last_status = None

        while True:
            time.sleep(current_interval)
            try:
                response = self._client.get(f"{self.base_api_url}/jobs/{job_id}/status")
                response.raise_for_status()
                data = response.json()
                status = data.get("status", "")
                if status == "done":
                    download = data.get("download", {})
                    download_url = download.get("download_url", "")
                    if not download_url:
                        raise RuntimeError(
                            "Failed to receive download URL from server."
                        )

                    if download_dir is None and filename is None:
                        # Return bytes content instead of saving to file
                        file_bytes = self.download_bytes(download_url)
                        return self.extract_chunks_and_chunksets_from_poma_archive(
                            poma_archive_data=file_bytes
                        )
                    else:
                        # Save downloaded file (to download_dir or current dir if only filename set)
                        save_filename = (
                            filename or download.get("filename") or f"{job_id}.poma"
                        )
                        if not save_filename.endswith(".poma"):
                            save_filename = f"{save_filename}.poma"
                        save_dir = (
                            download_dir
                            if download_dir not in (None, "")
                            else "."
                        )
                        downloaded_file_path = self.download_file(
                            download_url, save_filename, save_directory=save_dir
                        )
                        return self.extract_chunks_and_chunksets_from_poma_archive(
                            poma_archive_path=downloaded_file_path
                        )
                elif status == "failed":
                    error_code = data.get("code", "unknown")
                    error_details = data.get("error", "No details provided.")
                    error_message = (
                        f"Job failed with code {error_code}: {error_details}"
                    )
                    raise ApiError(
                        f"Job {job_id} failed: {data.get('error', error_message)}",
                        job_id=job_id,
                    )
                elif status == "processing":
                    if show_progress:
                        print(f"Job {job_id} is still processing...")
                    if last_status == "pending":
                        current_interval = poll_interval
                    current_interval = min(current_interval * 1.5, max_interval)
                elif status == "pending":
                    if show_progress:
                        print(
                            f"Job {job_id} is pending (queued due to rate limiting, sequential processing - common on demo accounts)..."
                        )
                    current_interval = min(current_interval * 1.5, max_interval)
                else:
                    raise InvalidResponseError(
                        f"Unexpected job status for job {job_id}: {status}"
                    )
                last_status = status
            except httpx.HTTPStatusError as error:
                status = error.response.status_code
                body = (error.response.text or "").strip()
                msg = f"Failed to get job status for '{job_id}'"
                if body:
                    msg = f"{msg}: {body}"
                raise ApiError(
                    msg,
                    status_code=status,
                    response_body=body,
                    job_id=job_id,
                ) from error
            except PomaSDKError:
                raise
            except Exception as error:
                raise RuntimeError(f"POMA-AI job polling failed: {error}") from error

    def extract_chunks_and_chunksets_from_poma_archive(
        self,
        poma_archive_data: bytes | None = None,
        poma_archive_path: str | os.PathLike[str] | None = None,
    ) -> dict[str, Any]:
        """
        Extract chunks and chunksets from a POMA archive; delegates to module-level function.
        POMA archive is a zip file containing chunks.json and chunksets.json.
        Args:
            poma_archive_data: The POMA archive as bytes.
            poma_archive_path: Path to the POMA archive file.
        Returns:
            dict: A dictionary with ``chunks`` and ``chunksets`` keys.
        """
        return extract_chunks_and_chunksets_from_poma_archive(
            poma_archive_data=poma_archive_data,
            poma_archive_path=poma_archive_path,
        )

    def create_cheatsheet(
        self,
        relevant_chunksets: list[dict[str, Any]],
        all_chunks: list[dict[str, Any]],
    ) -> str:
        """
        Generates a single cheatsheet for one single document
        from relevant chunksets (relevant for a certain query)
        and from all chunks of that document (providing the textual content).

        Args:
            relevant_chunksets (list[dict]): A list of chunksets, each containing a "chunks" key with a list of chunk IDs.
            all_chunks (list[dict]): A list of all chunk dictionaries of the same document, each representing a chunk of content.
        Returns:
            str: The textual content of the generated cheatsheet.

        Raises:
            InvalidInputError: If ``relevant_chunksets`` or ``all_chunks`` is not provided (None),
                not a list, or empty.
            InvalidResponseError: If the generated result has an unexpected shape (e.g. no content).
            ValueError: If the underlying generator raises (e.g. invalid chunkset structure,
                duplicate chunk_index, or file_id mismatch).
        """
        if relevant_chunksets is None or all_chunks is None:
            raise InvalidInputError(
                "relevant_chunksets and all_chunks must be provided (cannot be None)."
            )
        if not isinstance(relevant_chunksets, list) or not isinstance(all_chunks, list):
            raise InvalidInputError(
                "relevant_chunksets and all_chunks must be lists."
            )
        if not relevant_chunksets or not all_chunks:
            raise InvalidInputError(
                "relevant_chunksets and all_chunks must be non-empty."
            )
        cheatsheets = generate_cheatsheets(relevant_chunksets, all_chunks)
        if (
            not cheatsheets
            or not isinstance(cheatsheets, list)
            or len(cheatsheets) == 0
            or "content" not in cheatsheets[0]
        ):
            raise InvalidResponseError(
                "Cheatsheet could not be created from input chunks (unexpected result shape)."
            )
        return cheatsheets[0]["content"]

    def create_cheatsheets(
        self,
        relevant_chunksets: list[dict[str, Any]],
        all_chunks: list[dict[str, Any]],
    ) -> list[dict[str, Any]]:
        """
        Generates cheatsheets from relevant chunksets (relevant for a certain query)
        and from all the chunks of all affected documents (providing the textual content).
        One cheatsheet is created for each document found in the chunks (tagged with file_id).

        Args:
            relevant_chunksets (list[dict]): A list of chunksets, each containing a "chunks" key with a list of chunk IDs.
            all_chunks (list[dict]): A list of all available chunk dictionaries of affected documents, each representing a chunk of content.
        Returns:
            list[dict]: A list of dictionaries representing the generated cheatsheets, each containing:
                - 'file_id': The tag associated with the respective document.
                - 'content': The textual content of the generated cheatsheet.

        Raises:
            InvalidInputError: If ``relevant_chunksets`` or ``all_chunks`` is not provided (None) or not a list.
            ValueError: If the underlying generator raises (e.g. invalid chunkset structure,
                duplicate chunk_index, or file_id mismatch).
        """
        if relevant_chunksets is None or all_chunks is None:
            raise InvalidInputError(
                "relevant_chunksets and all_chunks must be provided (cannot be None)."
            )
        if not isinstance(relevant_chunksets, list) or not isinstance(all_chunks, list):
            raise InvalidInputError(
                "relevant_chunksets and all_chunks must be lists."
            )
        return generate_cheatsheets(relevant_chunksets, all_chunks)

    def download_file(
        self,
        download_url: str,
        filename: str | None = None,
        *,
        save_directory: str | os.PathLike[str] | None = None,
    ) -> str:
        """
        Download a file from the given download URL.
        Args:
            download_url (str):
                The URL to download the file from.
            filename (str, optional):
                The filename to save the file as. If not provided, will be extracted from URL.
            save_directory (str | os.PathLike[str], optional):
                Directory to save the file in. If not provided, saves to current directory.
        Returns:
            str: The path to the downloaded file.
        """
        if not download_url:
            raise ValueError("download_url cannot be empty")

        # Determine filename
        if not filename:
            filename = Path(download_url).name or "downloaded_file"

        # Determine save directory (default current directory so path has a parent)
        if save_directory:
            save_path = Path(save_directory) / filename
        else:
            save_path = Path(".") / filename

        # Create the directory if it doesn't exist (skip when file is in cwd)
        parent = os.path.dirname(save_path)
        if parent:
            os.makedirs(parent, exist_ok=True)

        # Download the file data
        content = self.download_bytes(download_url)

        # Save the file
        with open(save_path, "wb") as f:
            f.write(content)

        return str(save_path)

    def download_bytes(
        self,
        download_url: str,
    ) -> bytes:
        """
        Download a file from the given download URL and return the bytes content.

        On HTTP error, raises :exc:`ApiError` with :attr:`ApiError.status_code` and
        :attr:`ApiError.response_body` set from the API. For job downloads
        (``GET /jobs/{job_id}/download``), typical codes: 400 (invalid job_id/UUID),
        401 (authorization), 403 (account not found), 404 (job download not found),
        500 (failed to download job data).

        Args:
            download_url (str):
                The URL to download the file from.
        Returns:
            bytes: The content of the downloaded file as bytes.

        Raises:
            ApiError: On API HTTP error (see :attr:`ApiError.status_code` and
                :attr:`ApiError.response_body`).
        """
        if not download_url:
            raise ValueError("download_url cannot be empty")

        # Construct the full URL if it's a relative path
        if download_url.startswith("/"):
            full_url = f"{self.base_api_url}{download_url}"
        else:
            full_url = download_url

        print("Downloading file from:", full_url)
        try:
            response = self._client.get(full_url)
            response.raise_for_status()
            return response.content
        except httpx.HTTPStatusError as error:
            status = error.response.status_code
            body = (error.response.text or "").strip()
            msg = f"Failed to download '{download_url}'"
            if body:
                msg = f"{msg}: {body}"
            raise ApiError(
                msg,
                status_code=status,
                response_body=body,
                url=download_url,
            ) from error
        except PomaSDKError:
            raise
        except Exception as error:
            raise RuntimeError(f"File download failed: {error}") from error

    def close(self):
        self._client.close()

    def __enter__(self) -> "Poma":
        return self

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