from rag.agents.interface import Pipeline
from sparrow_parse.vllm.inference_factory import InferenceFactory
from sparrow_parse.extractors.vllm_extractor import VLLMExtractor
import timeit
from rich import print
from rich.progress import Progress, SpinnerColumn, TextColumn
from typing import Any, List
from .sparrow_validator import Validator
from .sparrow_utils import is_valid_json, get_json_keys_as_string
import warnings
import os


warnings.filterwarnings("ignore", category=DeprecationWarning)
warnings.filterwarnings("ignore", category=UserWarning)


class SparrowParsePipeline(Pipeline):

    def __init__(self):
        pass

    def run_pipeline(self,
                     payload: str,
                     query_inputs: [str],
                     query_types: [str],
                     keywords: [str],
                     query: str,
                     file_path: str,
                     index_name: str,
                     options: List[str] = None,
                     group_by_rows: bool = True,
                     update_targets: bool = True,
                     debug: bool = False,
                     local: bool = True) -> Any:
        print(f"\nRunning pipeline with {payload}\n")

        start = timeit.default_timer()

        query_all_data = False
        if query == "*":
            query_all_data = True
            query = None
        else:
            try:
                query, query_schema = self.invoke_pipeline_step(lambda: self.prepare_query_and_schema(query, debug),
                                                          "Preparing query and schema", local)
            except ValueError as e:
                raise e

        llm_output = self.invoke_pipeline_step(lambda: self.execute_query(options, query_all_data, query, file_path, debug),
                                               "Executing query", local)

        validation_result = None
        if query_all_data is False:
            validation_result = self.invoke_pipeline_step(lambda: self.validate_result(llm_output, query_all_data, query_schema, debug),
                                                      "Validating result", local)

        end = timeit.default_timer()

        print(f"Time to retrieve answer: {end - start}")

        return validation_result if validation_result is not None else llm_output


    def prepare_query_and_schema(self, query, debug):
        is_query_valid = is_valid_json(query)
        if not is_query_valid:
            raise ValueError("Invalid query. Please provide a valid JSON query.")

        query_keys = get_json_keys_as_string(query)
        query_schema = query
        query = "retrieve " + query_keys

        query = query + ". return response in JSON format, by strictly following this JSON schema: " + query_schema

        return query, query_schema


    def execute_query(self, options, query_all_data, query, file_path, debug):
        extractor = VLLMExtractor()

        # export HF_TOKEN="hf_"
        config = {}
        if options[0] == 'huggingface':
            config = {
                "method": options[0],  # Could be 'huggingface' or 'local_gpu'
                "hf_space": options[1],
                "hf_token": os.getenv('HF_TOKEN')
            }
        else:
            # Handle other cases if needed
            return "First element is not 'huggingface'"

        # Use the factory to get the correct instance
        factory = InferenceFactory(config)
        model_inference_instance = factory.get_inference_instance()

        input_data = [
            {
                "image": file_path,
                "text_input": query
            }
        ]

        # Now you can run inference without knowing which implementation is used
        llm_output = extractor.run_inference(model_inference_instance, input_data, generic_query=query_all_data,
                                             debug=debug)

        return llm_output


    def validate_result(self, llm_output, query_all_data, query_schema, debug):
        validator = Validator(query_schema)

        validation_result = validator.validate_json_against_schema(llm_output, validator.generated_schema)
        if validation_result is not None:
            return validation_result
        else:
            if debug:
                print("LLM output is valid according to the schema.")


    def invoke_pipeline_step(self, task_call, task_description, local):
        if local:
            with Progress(
                    SpinnerColumn(),
                    TextColumn("[progress.description]{task.description}"),
                    transient=False,
            ) as progress:
                progress.add_task(description=task_description, total=None)
                ret = task_call()
        else:
            print(task_description)
            ret = task_call()

        return ret