import os
import asyncio
import shutil
import tempfile
from pathlib import Path
from typing import Dict, List, Optional, Type, Union

from .types import VlenseResponse
from ..lib.retrieval import (
    get_collection_path,
    load_manifest,
    save_manifest,
    resolve_input_files,
    parse_indexed_pages,
    materialize_document_pages,
    rank_pages,
    get_embeddings_path,
)

from pydantic import BaseModel  # Import BaseModel for type hinting


class Vlense:
    """
    Vlense class for performing OCR and handling document queries.
    """

    def __init__(self):
        pass

    async def ocr(
        self,
        file_path: Union[str, List[str]],
        model: str = "gemini-flash-latest",
        output_dir: Optional[str] = None,
        temp_dir: Optional[str] = None,
        batch_size: int = 3,
        format: str = "markdown",
        json_schema: Optional[Type[BaseModel]] = None,  # New parameter
        clean_temp_files: Optional[bool] = True,
    ) -> Dict[str, VlenseResponse]:
        """
        API to perform OCR to markdown, html, or structured JSON using Vision models.
        Supports both PDF and image files.
        Please set up the environment variables for the model and model provider before using this API. Refer: https://docs.litellm.ai/docs/providers

        :param file_path: The path or URL to the PDF/image file to process.
        :type file_path: Union[str, List[str]]
        :param model: The model to use for generating completions, defaults to "gemini-flash-latest". Note - Refer: https://docs.litellm.ai/docs/providers to pass the correct model name as per provider specifications.
        :type model: str, optional
        :param output_dir: The directory to save the output, defaults to None.
        :type output_dir: Optional[str], optional
        :param temp_dir: The directory to store temporary files, defaults to a named folder in the system's temp directory. If it exists, its contents will be deleted.
        :type temp_dir: Optional[str], optional
        :param batch_size: The number of concurrent processes to run, defaults to 3.
        :type batch_size: int, optional
        :param format: Output format, either 'markdown', 'html', or 'json'.
        :type format: str, optional
        :param json_schema: Pydantic model defining the JSON schema for structured output, required if format is 'json'.
        :type json_schema: Optional[Type[BaseModel]], optional
        :param clean_temp_files: Whether to cleanup the temporary files after processing, defaults to True.
        :type clean_temp_files: Optional[bool], optional
        :return: The content generated by the model.
        :rtype: Dict[str, VlenseResponse]
        """

        if format not in ["markdown", "html", "json"]:
            raise ValueError("Format must be either 'markdown', 'html', or 'json'")

        if format == "json" and not json_schema:
            raise ValueError("json_schema must be provided when format is 'json'")

        results = {}

        if isinstance(file_path, str):
            file_paths = [file_path]
        else:
            file_paths = file_path

        if output_dir:
            os.makedirs(output_dir, exist_ok=True)

        if temp_dir:
            if os.path.exists(temp_dir):
                await asyncio.to_thread(shutil.rmtree, temp_dir)
            os.makedirs(temp_dir, exist_ok=True)

        from ..lib.pdf import process_batch_with_completion
        from ..models.litellmmodel import LiteLLMModel

        llm_model = LiteLLMModel(model=model, format=format, json_schema=json_schema)

        with tempfile.TemporaryDirectory() as temp_dir_:
            if temp_dir:
                temp_directory = temp_dir
            else:
                temp_directory = temp_dir_

            # Process all files using the updated process_batch_with_completion
            batch_results = await process_batch_with_completion(
                model=llm_model,
                file_paths=file_paths,
                batch_size=batch_size,
                format=format,
                output_dir=output_dir,
                temp_directory=temp_directory,
            )

            results.update(batch_results)

            if clean_temp_files and temp_dir:
                shutil.rmtree(temp_dir)

        return results

    async def index(
        self,
        data_dir: Union[str, List[str]],
        collection_name: str,
        index_dir: str = ".vlense",
        retriever_model: str = "vidore/colSmol-500M",
        temp_dir: Optional[str] = None,
        embedding_batch_size: int = 2,
    ) -> str:
        """
        Index PDFs or images with ColPali-engine retrieval for later QA.
        """
        file_paths = resolve_input_files(data_dir)
        return await self._index_with_colpali(
            file_paths=file_paths,
            collection_name=collection_name,
            index_dir=index_dir,
            retriever_model=retriever_model,
            temp_dir=temp_dir,
            embedding_batch_size=embedding_batch_size,
        )

    async def ask(
        self,
        query: str,
        collection_name: str,
        model: str = "gemini-flash-latest",
        index_dir: str = ".vlense",
        top_k: int = 3,
    ) -> str:
        """
        Answer a question using colpali-engine retrieval over an indexed collection.
        """
        manifest = load_manifest(index_dir=index_dir, collection_name=collection_name)
        return await self._ask_with_colpali(
            query=query,
            collection_name=collection_name,
            model=model,
            index_dir=index_dir,
            top_k=top_k,
            retriever_model=manifest.get("retriever_model", "vidore/colSmol-500M"),
        )

    async def _index_with_colpali(
        self,
        file_paths: List[str],
        collection_name: str,
        index_dir: str,
        retriever_model: str,
        temp_dir: Optional[str],
        embedding_batch_size: int,
    ) -> str:
        from ..models.colpali import ColPaliRetriever

        collection_path = get_collection_path(index_dir, collection_name)
        if collection_path.exists():
            shutil.rmtree(collection_path)
        pages_root = collection_path / "pages"
        pages_root.mkdir(parents=True, exist_ok=True)

        with tempfile.TemporaryDirectory() as temp_dir_:
            working_temp_dir = temp_dir or temp_dir_
            if temp_dir:
                if os.path.exists(temp_dir):
                    await asyncio.to_thread(shutil.rmtree, temp_dir)
                os.makedirs(temp_dir, exist_ok=True)

            page_entries = []
            page_image_paths = []
            for file_path in file_paths:
                entries = await materialize_document_pages(
                    source_path=file_path,
                    pages_root=pages_root,
                    temp_directory=working_temp_dir,
                )
                page_entries.extend(entries)
                page_image_paths.extend(entry["image_path"] for entry in entries)

            if temp_dir:
                shutil.rmtree(temp_dir)

        retriever = ColPaliRetriever(model_name=retriever_model)
        page_embeddings = retriever.encode_images(
            page_image_paths,
            batch_size=embedding_batch_size,
        )
        retriever.save_embeddings(
            page_embeddings,
            str(get_embeddings_path(index_dir, collection_name)),
        )

        return save_manifest(
            index_dir=index_dir,
            collection_name=collection_name,
            retriever_model=retriever_model,
            pages=parse_indexed_pages({"pages": page_entries}),
        )

    async def _ask_with_colpali(
        self,
        query: str,
        collection_name: str,
        model: str,
        index_dir: str,
        top_k: int,
        retriever_model: str,
    ) -> str:
        from ..models.colpali import ColPaliRetriever
        from ..models.litellmmodel import LiteLLMModel

        manifest = load_manifest(index_dir=index_dir, collection_name=collection_name)
        pages = parse_indexed_pages(manifest)
        if not pages:
            raise ValueError(f"Collection '{collection_name}' does not contain any indexed pages.")

        retriever = ColPaliRetriever(model_name=retriever_model)
        page_embeddings = retriever.load_embeddings(manifest["embeddings_path"])
        query_embeddings = retriever.encode_queries([query])
        scores = retriever.score(
            query_embeddings=query_embeddings,
            document_embeddings=page_embeddings,
        )
        results = rank_pages(
            pages=pages,
            scores=scores,
            top_k=max(1, min(top_k, len(pages))),
        )
        if not results:
            raise ValueError(f"Collection '{collection_name}' returned no matches for the query.")

        image_paths = []
        page_references = []
        for result in results:
            image_paths.append(result.page.image_path)
            page_references.append(f"{result.page.file_name} p.{result.page.page_number}")

        llm_model = LiteLLMModel(model=model, format="markdown")
        answer = await llm_model.answer_question(
            question=query,
            image_paths=image_paths,
            page_references=page_references,
        )
        return answer.content
