#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# ================================================== #
# This file is a part of PYGPT package               #
# Website: https://pygpt.net                         #
# GitHub:  https://github.com/szczyglis-dev/py-gpt   #
# MIT License                                        #
# Created By  : Marcin Szczygliński                  #
# Updated Date: 2026.01.06 06:00:00                  #
# ================================================== #

import os
import time
import mimetypes
from typing import Optional, List, Dict, Any, Union

from pygpt_net.item.store import RemoteStoreItem

from .worker.importer import Importer

class Store:
    def __init__(self, window=None):
        """
        Google (Gemini) File Search stores API wrapper

        :param window: Window instance
        """
        self.window = window
        self.importer = Importer(window)

    def get_client(self):
        """
        Get Google GenAI client

        :return: google.genai.Client
        """
        return self.window.core.api.google.get_client()

    def log(self, msg: str, callback: Optional[callable] = None):
        """
        Log message

        :param msg: message to log
        :param callback: callback log function
        """
        if callback is not None:
            callback(msg)
        else:
            print(msg)

    # -----------------------------
    # helpers
    # -----------------------------

    def _clamp(self, val: int, lo: int, hi: int) -> int:
        """Clamp integer value into range [lo, hi]."""
        try:
            v = int(val or 0)
        except Exception:
            v = lo
        if v < lo:
            v = lo
        if v > hi:
            v = hi
        return v

    def _download_dir(self) -> str:
        """
        Resolve target download directory (uses download.dir if set).
        """
        if self.window.core.config.has("download.dir") and self.window.core.config.get("download.dir") != "":
            dir_path = os.path.join(
                self.window.core.config.get_user_dir('data'),
                self.window.core.config.get("download.dir"),
            )
        else:
            dir_path = self.window.core.config.get_user_dir('data')
        os.makedirs(dir_path, exist_ok=True)
        return dir_path

    def _ensure_unique_path(self, dir_path: str, filename: str) -> str:
        """
        Ensure unique filename in dir, add timestamp prefix if exists.
        """
        path = os.path.join(dir_path, filename)
        if os.path.exists(path):
            prefix = time.strftime("%Y%m%d_%H%M%S_")
            path = os.path.join(dir_path, f"{prefix}{filename}")
        return path

    def _guess_filename(self, file_meta: Any, fallback: str = "downloaded.bin") -> str:
        """
        Best-effort filename from File metadata or URI.
        """
        name = None
        for attr in ("display_name", "filename", "name", "file_name"):
            try:
                val = getattr(file_meta, attr, None)
                if not name and isinstance(val, str) and val:
                    name = os.path.basename(val)
            except Exception:
                pass
            if not name and isinstance(file_meta, dict):
                val = file_meta.get(attr)
                if isinstance(val, str) and val:
                    name = os.path.basename(val)

        if not name:
            # allow URI-like strings
            if isinstance(file_meta, str):
                name = os.path.basename(file_meta.split("?")[0].split("#")[0])
        return name or fallback

    # -----------------------------
    # Files service (global)
    # -----------------------------

    def get_file(self, file_name: str):
        """
        Get Files API file metadata by name (e.g. 'files/abc-123').

        :param file_name: Files API resource name
        :return: file metadata
        """
        client = self.get_client()
        return client.files.get(name=file_name)

    def upload(self, path: str) -> Optional[str]:
        """
        Upload file to Files API (not to a File Search Store). Use when you want to import later.

        :param path: file path
        :return: Files API resource name (e.g., 'files/abc-123') or None
        """
        client = self.get_client()
        if not os.path.exists(path):
            return None

        cfg = {'name': os.path.basename(path)}
        res = client.files.upload(file=path, config=cfg)
        if res is not None and hasattr(res, "name"):
            return res.name

    def delete_file(self, file_name: str) -> Optional[str]:
        """
        Delete file from Files API

        :param file_name: Files API resource name ('files/...').
        :return: file_name if removed
        """
        client = self.get_client()
        res = client.files.delete(name=file_name)
        if res is not None:
            return file_name

    def download(self, file: Union[str, Any], path: str) -> bool:
        """
        Download a Files API item into the given path.

        :param file: file name ('files/...'), file object, or file URI
        :param path: target local path
        :return: True on success
        """
        client = self.get_client()
        data = None
        try:
            data = client.files.download(file=file)
        except Exception:
            pass
        if not data:
            return False
        # google-genai returns bytes
        try:
            with open(path, "wb") as f:
                f.write(data if isinstance(data, (bytes, bytearray)) else bytes(data))
            return True
        except Exception:
            return False

    def download_to_dir(self, file: Union[str, Any], prefer_name: Optional[str] = None) -> Optional[str]:
        """
        Download a Files API item into configured download directory.

        :param file: file name ('files/...'), file object, or file URI
        :param prefer_name: optional preferred filename
        :return: saved path or None
        """
        dir_path = self._download_dir()
        filename = None

        # Try to resolve filename from metadata
        file_meta = None
        try:
            name = None
            if isinstance(file, str) and file.startswith("files/"):
                name = file
            elif hasattr(file, "name"):
                name = getattr(file, "name", None)

            if name:
                file_meta = self.get_file(name)
        except Exception:
            file_meta = None

        if prefer_name and isinstance(prefer_name, str):
            filename = os.path.basename(prefer_name)

        if not filename:
            filename = self._guess_filename(file_meta if file_meta is not None else file)

        # Infer extension from mime, if missing
        if not os.path.splitext(filename)[1] and file_meta is not None:
            try:
                mime = getattr(file_meta, "mime_type", None)
                if isinstance(file_meta, dict):
                    mime = file_meta.get("mime_type", mime)
                if mime:
                    ext = mimetypes.guess_extension(mime) or ""
                    if ext:
                        filename = filename + ext
            except Exception:
                pass

        path = self._ensure_unique_path(dir_path, filename)
        if self.download(file, path):
            return path
        return None

    def get_files_ids_all(
            self,
            items: list,
            order: str = "asc",
            limit: int = 100,
            after: Optional[str] = None,
    ) -> list:
        """
        Get all Files API file names (paginated)

        :param items: items accumulator
        :param order: unused, for compatibility
        :param limit: page size (max 100 for Files API)
        :param after: page token
        :return: file names list
        """
        client = self.get_client()
        limit = self._clamp(limit, 1, 100)  # Files API allows up to 100
        cfg: Dict[str, Any] = {'page_size': limit}
        if after:
            cfg['page_token'] = after

        pager = client.files.list(config=cfg)
        for f in pager:
            name = getattr(f, "name", None)
            if name and name not in items:
                items.append(name)

        # The SDK pager auto-iterates pages; explicit tokens are optional
        return items

    def get_files_ids(self) -> List[str]:
        """
        Get Files API file names (single iterator over all pages)

        :return: list of file names
        """
        client = self.get_client()
        items = []
        pager = client.files.list()
        for f in pager:
            name = getattr(f, "name", None)
            if name and name not in items:
                items.append(name)
        return items

    def remove_files(self, callback: Optional[callable] = None) -> int:
        """
        Remove all Files API files

        :param callback: callback function
        :return: number of deleted files
        """
        num = 0
        files = self.get_files_ids()
        for file_name in files:
            self.log("Removing file: " + file_name, callback)
            try:
                res = self.delete_file(file_name)
                if res:
                    num += 1
            except Exception as e:
                msg = "Error removing file {}: {}".format(file_name, str(e))
                self.log(msg, callback)
        return num

    def remove_file(self, file_name: str, callback: Optional[callable] = None) -> bool:
        """
        Remove a single Files API file

        :param file_name: Files API resource name ('files/...').
        :param callback: callback function
        :return: True if removed
        """
        self.log("Removing file: " + file_name, callback)
        try:
            res = self.delete_file(file_name)
            return res is not None
        except Exception as e:
            msg = "Error removing file {}: {}".format(file_name, str(e))
            self.log(msg, callback)
            raise

    # -----------------------------
    # File Search stores
    # -----------------------------

    def import_stores(
            self,
            items: dict,
            order: str = "asc",
            limit: int = 100,
            after: Optional[str] = None,
            callback: Optional[callable] = None
    ) -> dict:
        """
        Import File Search stores

        :param items: items dict accumulator
        :param order: unused
        :param limit: page size (Google limit: 20)
        :param after: page token
        :param callback: callback
        :return: items dict of RemoteStoreItem
        """
        client = self.get_client()
        limit = self._clamp(limit, 1, 20)  # API limit = 20
        cfg: Dict[str, Any] = {'page_size': limit}
        if after:
            cfg['page_token'] = after

        pager = client.file_search_stores.list(config=cfg)
        for remote in pager:
            name = getattr(remote, "name", None)  # 'fileSearchStores/...'
            if not name:
                continue
            if name not in items:
                items[name] = RemoteStoreItem()
            items[name].id = name
            items[name].name = getattr(remote, "display_name", "") or ""
            items[name].file_ids = []
            items[name].provider = "google"

            status = self.window.core.remote_store.google.parse_status(remote)
            self.window.core.remote_store.google.append_status(items[name], status)
            self.log("Imported file search store: " + name, callback)

        return items

    def create_store(self, name: str, expire_days: int = 0):
        """
        Create File Search store

        :param name: display name
        :param expire_days: ignored (not supported)
        :return: store object
        """
        client = self.get_client()
        store = client.file_search_stores.create(config={'display_name': name})
        if store is not None:
            return store

    def update_store(self, id: str, name: str, expire_days: int = 0):
        """
        Update File Search store (no remote update available). Returns the current remote store.

        :param id: store name ('fileSearchStores/...').
        :param name: local alias (persisted locally)
        :param expire_days: ignored
        :return: store object (fresh get) or None if not found
        """
        try:
            return self.get_store(id)
        except Exception:
            return None

    def get_store(self, id: str):
        """
        Get File Search store

        :param id: store name ('fileSearchStores/...').
        :return: store object
        """
        client = self.get_client()
        return client.file_search_stores.get(name=id)

    def remove_store(self, id: str):
        """
        Delete File Search store (force delete related documents)

        :param id: store name ('fileSearchStores/...').
        :return: delete response
        """
        client = self.get_client()
        return client.file_search_stores.delete(name=id, config={'force': True})

    def get_stores_ids(
            self,
            items: list,
            order: str = "asc",
            limit: int = 100,
            after: Optional[str] = None,
    ) -> list:
        """
        Get all File Search store names

        :param items: items accumulator
        :param order: unused
        :param limit: page size (Google limit 20)
        :param after: page token
        :return: list of store names
        """
        client = self.get_client()
        limit = self._clamp(limit, 1, 20)
        cfg: Dict[str, Any] = {'page_size': limit}
        if after:
            cfg['page_token'] = after

        pager = client.file_search_stores.list(config=cfg)
        for remote in pager:
            name = getattr(remote, "name", None)
            if name and name not in items:
                items.append(name)
        return items

    # -----------------------------
    # Documents within a store
    # -----------------------------

    def get_store_files_ids(
            self,
            store_id: str,
            items: list,
            order: str = "asc",
            limit: int = 100,
            after: Optional[str] = None,
    ) -> list:
        """
        Get all Document names for a File Search store

        :param store_id: store name ('fileSearchStores/...').
        :param items: items accumulator
        :param order: unused
        :param limit: page size (Google limit 20)
        :param after: page token
        :return: list of Document names ('fileSearchStores/.../documents/...').
        """
        client = self.get_client()
        limit = self._clamp(limit, 1, 20)
        cfg: Dict[str, Any] = {'page_size': limit}
        if after:
            cfg['page_token'] = after

        try:
            pager = client.file_search_stores.documents.list(parent=store_id, config=cfg)
        except Exception:
            pager = client.documents.list(parent=store_id, config=cfg)

        for remote in pager:
            name = getattr(remote, "name", None)
            if name and name not in items:
                items.append(name)
        return items

    def remove_from_stores(self) -> int:
        """
        Remove all documents from all File Search stores

        :return: number of deleted documents
        """
        stores = self.get_stores_ids([])
        num = 0
        for store_id in stores:
            files = self.get_store_files_ids(store_id, [])
            for doc_name in files:
                self.log("Removing document from store [{}]:{} ".format(store_id, doc_name))
                self.delete_store_file(store_id, doc_name)
                num += 1
        return num

    def remove_from_store(self, store_id: str) -> int:
        """
        Remove all documents from a specific File Search store

        :param store_id: store name ('fileSearchStores/...').
        :return: number of deleted documents
        """
        files = self.get_store_files_ids(store_id, [])
        num = 0
        for doc_name in files:
            self.log("Removing document from store [{}]:{} ".format(store_id, doc_name))
            self.delete_store_file(store_id, doc_name)
            num += 1        
        return num

    def remove_all(self, callback: Optional[callable] = None) -> int:
        """
        Remove all File Search stores

        :param callback: callback function
        :return: number of deleted stores
        """
        num = 0
        stores = self.get_stores_ids([])
        for store_id in stores:
            self.log("Removing file search store: " + store_id, callback)
            try:
                self.remove_store(store_id)
                num += 1
            except Exception as e:
                msg = "Error removing file search store {}: {}".format(store_id, str(e))
                self.log(msg, callback)
        return num

    def add_file(self, store_id: str, file_name: str):
        """
        Import a Files API file into a File Search store.

        :param store_id: store name ('fileSearchStores/...').
        :param file_name: Files API file name ('files/...').
        :return: operation (long running)
        """
        client = self.get_client()
        op = client.file_search_stores.import_file(
            file_search_store_name=store_id,
            file_name=file_name,
        )
        return op

    def delete_store_file(self, store_id: str, document_name: str):
        """
        Delete a Document from a File Search store.

        :param store_id: store name ('fileSearchStores/...').
        :param document_name: document resource name ('fileSearchStores/.../documents/...').
        :return: delete response
        """
        client = self.get_client()
        return client.file_search_stores.documents.delete(name=document_name, config={'force': True})

    def remove_store_file(self, store_id: str, document_name: str):
        """
        Compatibility alias used by controller: remove a document from a File Search store.
        """
        return self.delete_store_file(store_id, document_name)

    def import_stores_files(self, callback: Optional[callable] = None) -> int:
        """
        Import all documents (files) from all File Search stores

        :param callback: callback function
        :return: number of imported documents
        """
        store_ids = self.get_stores_ids([])
        num = 0
        for store_id in store_ids:
            items = []
            try:
                items = self.import_store_files(store_id, items, callback=callback)
            except Exception as e:
                msg = "Error importing store {} documents list: {}".format(store_id, str(e))
                self.log(msg, callback)
            num += len(items)
        return num

    def import_store_files(
            self,
            store_id: str,
            items: list,
            order: str = "asc",
            limit: int = 100,
            after: Optional[str] = None,
            callback: Optional[callable] = None
    ) -> list:
        """
        Import a store's documents and insert into local DB

        :param store_id: store name ('fileSearchStores/...').
        :param items: accumulator
        :param order: unused
        :param limit: page size (Google limit 20)
        :param after: page token
        :param callback: log callback
        :return: list of imported document names
        """
        client = self.get_client()
        limit = self._clamp(limit, 1, 20)
        cfg: Dict[str, Any] = {'page_size': limit}
        if after:
            cfg['page_token'] = after

        try:
            pager = client.file_search_stores.documents.list(parent=store_id, config=cfg)
        except Exception:
            pager = client.documents.list(parent=store_id, config=cfg)

        for remote in pager:
            try:
                doc_name = getattr(remote, "name", None)
                if not doc_name:
                    continue
                if doc_name not in items:
                    items.append(doc_name)
                    try:
                        doc = client.file_search_stores.documents.get(name=doc_name)
                    except Exception:
                        doc = remote
                    self.window.core.remote_store.google.files.insert(store_id, doc)
                    msg = "Imported document {} to store {}".format(doc_name, store_id)
                    self.log(msg, callback)
            except Exception as e:
                msg = "Error importing document {} to store {}: {}".format(
                    getattr(remote, "name", "?"), store_id, str(e)
                )
                self.log(msg, callback)
        return items

    # -----------------------------
    # Convenience: direct upload to a store
    # -----------------------------

    def upload_to_store(self, store_id: str, path: str):
        """
        Directly upload a local file to a File Search store. This creates a Document.

        :param store_id: store name ('fileSearchStores/...').
        :param path: file path
        :return: created document object or None
        """
        client = self.get_client()
        if not os.path.exists(path):
            return None

        basename = os.path.basename(path)
        op = client.file_search_stores.upload_to_file_search_store(
            file=path,
            file_search_store_name=store_id,
            config={'display_name': basename},
        )

        for _ in range(180):
            if getattr(op, "done", False):
                break
            time.sleep(2)
            op = client.operations.get(op)

        try:
            docs = self.get_store_files_ids(store_id, [])
            for doc_name in docs:
                try:
                    d = client.file_search_stores.documents.get(name=doc_name)
                except Exception:
                    continue
                if getattr(d, "display_name", None) == basename:
                    return d
        except Exception:
            pass

        return None