import pandas as pd
import json
from pydantic import ValidationError
from werkzeug.datastructures import FileStorage

import dhuolib.utils as ut
from dhuolib.config import logger
from dhuolib.services import ServiceAPIMLFacade
from dhuolib.utils import validate_name
from dhuolib.validations import (
    ExperimentBody,
    PredictBatchModelBody,
    PredictOnlineModelBody,
    RunExperiment,
    TransitionModel,
    RunSearch,
    ModelBody
)


class DhuolibExperimentMLClient:
    def __init__(self, service_endpoint=None):
        if not service_endpoint:
            raise ValueError("service_endpoint is required")

        self.service = ServiceAPIMLFacade(service_endpoint)

    def create_experiment(
        self, experiment_name: str, experiment_tags: dict = None
    ) -> dict:
        try:
            experiment = ExperimentBody.model_validate({
                "experiment_name": validate_name(experiment_name),
                "experiment_tags": experiment_tags,
            })
        except ValidationError as e:
            logger.error(f"Error: {e}")
            return {"error": str(e)}

        response = self.service.experiment_api.create(experiment.model_dump())
        logger.info(
            f"Experiment Name: {experiment.experiment_name} created"
        )

        return response

    def predict_online(self, run_params) -> dict:
        try:
            PredictOnlineModelBody.model_validate(run_params)
        except ValidationError as e:
            logger.error(f"Error: {e}")
            return {"error": str(e)}
        try:
            with open(run_params["data"], "rb") as f1:
                files = {
                    "data": FileStorage(
                        stream=f1, filename="data.csv", content_type="csv"
                    )
                }
                response = self.service.experiment_api.predict_online(
                    params=run_params, files=files)
                logger.info(f"Model Name: {run_params['modelname']} predictions")
                return response
        except FileNotFoundError as e:
            logger.error(f"Error: {e}")
            return {"error": str(e)}

    def _get_model_to_prediction(self, predict: PredictBatchModelBody) -> dict:

        path_to_pickle = self.service.experiment_api.download_pickle(
            model_name=predict.modelname,
            model_stage=predict.stage,
            run_id=predict.run_id,
            experiment_name=predict.experiment_name,
            local_filename=predict.batch_model_dir,
            type_model=predict.type_model,
        )

        return ut.load_pickle_model(path_to_pickle)

    def prediction_batch_with_dataframe(
        self, batch_params: dict, df: pd.DataFrame
    ) -> dict:
        try:
            predict = PredictBatchModelBody.model_validate(batch_params)
        except ValidationError as e:
            logger.error(f"Error: {e}")
            return {"error": str(e)}

        model = self._get_model_to_prediction(predict)
        return model.predict(df)

    def search_experiments(self, filter_string: str = "", max_results: int = 10, page_token: str = '') -> dict:

        response = self.service.experiment_api.search(
            filter_string=filter_string, max_results=max_results, page_token=page_token
        )
        return response

    def execute_run_for_experiment(
        self,
        type_model: str,
        experiment_name: str,
        modelpkl_path: str,
        requirements_path: str,
        tags={}
    ) -> dict:
        params = {
            "type_model": type_model,
            "experiment_name": experiment_name,
            "tags": json.dumps(tags) if isinstance(tags, dict) else tags

        }
        try:
            RunExperiment.model_validate(params)
        except ValidationError as e:
            logger.error(f"Error: {e}")
            return {"error": str(e)}

        try:
            with open(requirements_path, "rb") as f1, open(
                modelpkl_path, "rb"
            ) as f2:
                requirements_file = FileStorage(
                    stream=f1,
                    filename="requirements.txt",
                    content_type="text/plain",
                )
                model_pkl_file = FileStorage(
                    stream=f2,
                    filename="model.pkl",
                    content_type="application/octet-stream",
                )

                request_files = {
                    "requirements_path": (requirements_file.filename, requirements_file.stream, requirements_file.content_type),
                    "modelpkl_path": (model_pkl_file.filename, model_pkl_file.stream, model_pkl_file.content_type)
                }

                response = self.service.runs_api.create(
                    run_params=params, files=request_files)
                logger.info(f"Experiment Name: {params['experiment_name']} running")
                return response
        except FileNotFoundError as e:
            logger.error(f"Error: {e}")
            return {"error": str(e)}

    def search_runs(self, filter_string: str = "", max_results: int = 10, page_token: str = "", experiment_name: str = "") -> dict:
        try:
            RunSearch.model_validate({
                "experiment_name": experiment_name,
                "filter_string": filter_string,
                "max_results": max_results,
                "page_token": page_token,
            })
        except ValidationError as e:
            logger.error(f"Error: {e}")
            return {"error": str(e)}
        try:
            response = self.service.runs_api.search(
                filter_string=filter_string, max_results=max_results, page_token=page_token, experiment_name=experiment_name
            )
            return response
        except ValueError as e:
            logger.error(f"Error: {e}")
            return {"error": str(e)}

    def create_model(self, modelname: str, stage: str, run_id: str, model_uri: str, tags: dict) -> dict:
        try:
            transition_model = ModelBody.model_validate({
                "modelname": modelname,
                "stage": stage,
                "run_id": run_id,
                "model_uri": model_uri,
                "tags": tags,
            })
            response = self.service.model_api.create(transition_model.model_dump())
            return response
        except ValidationError as e:
            logger.error(f"Error: {e}")
            return {"error": str(e)}

    def transition_model(self, model_name: str, version: str, stage: str) -> dict:
        try:
            TransitionModel.model_validate(
                {"model_name": model_name, "version": version, "stage": stage})
        except ValidationError as e:
            logger.error(f"Error: {e}")
            return {"error": str(e)}

        response = self.service.model_api.transition_model(
            model_name=model_name, version=version, stage=stage)
        return response
 
    def search_models(self, filter_string: str = "", max_results: int = 10, page_token: str = "") -> dict:
        response = self.service.model_api.search(
            filter_string=filter_string, max_results=max_results, page_token=page_token)
        return response
