import json
import os
from abc import ABC, abstractmethod

import requests
from langchain_community.vectorstores.chroma import Chroma
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.schema import Document
from errors.errors import BadRequestException

from utils.local import get_data_context_id, get_session
from models.context import Context as ContextModel
from config import config
from consts.consts import VectorMode
from utils.logging import logging


CONTEXT_MODE_LOCAL = "local"
CONTEXT_MODE_REMOTE = "remote"
CONTEXT_MODE_DB = "db"


class ContextAPI(ABC):
    @abstractmethod
    def load_dataset_context(self):
        raise NotImplementedError()

    @abstractmethod
    def save_dataset_context(self, dataset_context):
        raise NotImplementedError()

    @abstractmethod
    def clear_dataset_context(self, key=None):
        raise NotImplementedError()

    @abstractmethod
    def load_session_context(self, session_context_id):
        raise NotImplementedError()

    @abstractmethod
    def save_session_context(self, session_context_id, session_context):
        raise NotImplementedError()

    @abstractmethod
    def load_domain_context(self):
        raise NotImplementedError()

    @abstractmethod
    def save_domain_context(self, domain_context):
        raise NotImplementedError()

    @abstractmethod
    def clear_session_context(self, key=None):
        raise NotImplementedError()


class RemoteContext(ContextAPI):
    def __init__(self) -> None:
        self.session_id = os.getenv("SESSION_ID")
        self.org_id = int(os.getenv("ORG_ID") or 0)
        self.cluster_id = int(os.getenv("CLUSTER_ID") or 0)
        self.db_name = os.getenv("DB_NAME")
        self.auth0_token = os.getenv("AUTH0_TOKEN")

        if not (self.session_id and self.org_id and self.cluster_id and self.db_name and self.auth0_token):
            raise BadRequestException("all of session_id/org_id/cluster_id/db_name/auth0_token are required")

        # optional/default value
        self.base_addr = "http://edaservice.nightly-ms.svc.cluster.local:11000"
        # self.base_addr = "http://127.0.0.1:11001"
        self.domain_context_id = int(os.getenv("DOMAIN_CONTEXT_ID") or 0)

        self.__ensure_session_context()
        self.__ensure_data_context()

    def __get_headers(self):
        headers = {
            "Authorization": "Bearer {}".format(self.auth0_token),
        }
        return headers

    def __ensure_session_context(self):
        url = self.base_addr + f"/aipower/api/v1/context/orgs/{self.org_id}/contexts"
        payload = {
            "type": "session",
            "name": f"session-{self.org_id}",
            "session_id": self.session_id,
            "creator": "jupyter",
            "data": {},
        }
        headers = self.__get_headers()
        resp = requests.put(url, json=payload, headers=headers)
        resp.raise_for_status()
        resp_json = resp.json()
        self.session_context_id = resp_json["result"]["id"]

    def __ensure_data_context(self):
        url = self.base_addr + f"/aipower/api/v1/context/orgs/{self.org_id}/contexts"
        payload = {
            "type": "data",
            "name": f"data-{self.org_id}",
            "cluster_id": self.cluster_id,
            "db_name": self.db_name,
            "creator": "jupyter",
            "data": {},
        }
        headers = self.__get_headers()
        resp = requests.put(url, json=payload, headers=headers)
        resp.raise_for_status()
        resp_json = resp.json()
        self.data_context_id = resp_json["result"]["id"]

    def __load_context_by_id(self, org_id, context_id):
        url = self.base_addr + f"/aipower/api/v1/context/orgs/{org_id}/contexts?context_id={context_id}"
        headers = self.__get_headers()
        resp = requests.get(url, headers=headers)
        resp.raise_for_status()

        resp_json = resp.json()

        return resp_json["result"]["data"]

    def __update_context_by_id(self, org_id, context_id, data):
        url = self.base_addr + f"/aipower/api/v1/context/orgs/{org_id}/contexts/{context_id}"
        payload = {
            "data": data,
        }
        headers = self.__get_headers()
        resp = requests.put(url, json=payload, headers=headers)
        resp.raise_for_status()

        resp_json = resp.json()

        return resp_json["result"]

    def __delete_context_by_id(self, org_id, context_id, key):
        url = self.base_addr + f"/aipower/api/v1/context/orgs/{org_id}/contexts/{context_id}"
        payload = {
            "key": key,
            "all": key is None,
        }
        headers = self.__get_headers()
        resp = requests.delete(url, json=payload, headers=headers)
        resp.raise_for_status()

        resp_json = resp.json()

        return resp_json["result"]

    def load_dataset_context(self):
        return self.__load_context_by_id(self.org_id, self.data_context_id)

    def save_dataset_context(self, dataset_context):
        self.__update_context_by_id(self.org_id, self.data_context_id, dataset_context)

    def clear_dataset_context(self, key=None):
        pass

    def load_session_context(self, session_context_id):
        return self.__load_context_by_id(self.org_id, self.session_context_id)

    def save_session_context(self, session_context_id, session_context):
        return self.__update_context_by_id(self.org_id, self.session_context_id, session_context)

    def clear_session_context(self, key=None):
        return self.__delete_context_by_id(self.org_id, self.session_context_id, key)

    def load_domain_context(self):
        pass

    def save_domain_context(self, domain_context):
        pass


class LocalContext(ContextAPI):
    def __init__(self) -> None:
        self.session_context_path = "./session_context.json"
        self.data_context_path = "./data_context.json"
        self.domain_context_path = "./domain_context.json"

    def __update_context_local(self, path, data):
        with open(path, "w") as f:
            json.dump(data, f)

    def __load_context_local(self, path):
        try:
            with open(path) as f:
                return json.load(f)
        except FileNotFoundError:
            with open(path, "w") as f:
                f.write("{}")
            return {}

    def __delete_context_local(self, path, key):
        try:
            if key:
                with open(path) as f:
                    data = json.load(f)
                    data.pop(key, None)
                    f.seek(0)
                    f.write(json.dumps(data))
            else:
                os.remove(path)
        except FileNotFoundError:
            return {}

    def load_dataset_context(self):
        return self.__load_context_local(self.data_context_path)

    def save_dataset_context(self, dataset_context):
        return self.__update_context_local(self.data_context_path, dataset_context)

    def clear_dataset_context(self, key=None):
        return self.__delete_context_local(self.data_context_path, key)

    def load_session_context(self, session_context_id):
        return self.__load_context_local(self.session_context_path)

    def save_session_context(self, session_context_id, session_context):
        return self.__update_context_local(self.session_context_path, session_context)

    def load_domain_context(self):
        return self.__load_context_local(self.domain_context_path)

    def save_domain_context(self, domain_context):
        return self.__update_context_local(self.domain_context_path, domain_context)

    def clear_session_context(self, key=None):
        return self.__delete_context_local(self.session_context_path, key)


class DBContext(ContextAPI):
    def load_dataset_context(self):
        data_context_id = get_data_context_id()
        s = get_session()

        data_context = ContextModel.get_by_id(s, data_context_id)
        assert data_context is not None

        return data_context.load_data()

    def save_dataset_context(self, dataset_context):
        data_context_id = get_data_context_id()
        s = get_session()

        data_context = ContextModel.get_by_id(s, data_context_id)
        assert data_context is not None

        logging.info("save_dataset_context %s of data_context_id %s", str(dataset_context), data_context_id)
        data_context.update_data(dataset_context)

    def clear_dataset_context(self, key=None):
        data_context_id = get_data_context_id()
        s = get_session()

        data_context = ContextModel.get_by_id(s, data_context_id)
        assert data_context is not None

        data_context.update_data({})

    def load_session_context(self, session_context_id):
        s = get_session()

        session_context = ContextModel.get_by_id(s, session_context_id)
        assert session_context is not None

        return session_context.load_data()

    def save_session_context(self, session_context_id, data):
        s = get_session()

        session_context = ContextModel.get_by_id(s, session_context_id)
        assert session_context is not None

        session_context.update_data(data)

    def load_domain_context(self):
        return self.__load_context_local(self.domain_context_path)

    def save_domain_context(self, domain_context):
        return self.__update_context_local(self.domain_context_path, domain_context)

    def clear_session_context(self, key=None):
        return self.__delete_context_local(self.session_context_path, key)


class Context(ContextAPI):
    def __init__(self) -> None:
        self.refresh()

    def refresh(self) -> None:
        # required
        self.context_mode = os.getenv("CONTEXT_MODE") or CONTEXT_MODE_DB

        if self.context_mode == CONTEXT_MODE_LOCAL:
            self.__context = LocalContext()
        elif self.context_mode == CONTEXT_MODE_REMOTE:
            self.__context = RemoteContext()
        elif self.context_mode == CONTEXT_MODE_DB:
            self.__context = DBContext()
        else:
            raise BadRequestException("invalid context mode")

    def load_dataset_context(self):
        return self.__context.load_dataset_context()

    def save_dataset_context(self, dataset_context):
        return self.__context.save_dataset_context(dataset_context)

    def clear_dataset_context(self, key=None):
        return self.__context.clear_dataset_context(key)

    def load_domain_context(self):
        return self.__context.load_domain_context()

    def save_domain_context(self, domain_context):
        return self.__context.save_domain_context(domain_context)

    def load_session_context(self, session_context_id):
        return self.__context.load_session_context(session_context_id)

    def save_session_context(self, session_context_id, session_context):
        return self.__context.save_session_context(session_context_id, session_context)

    def clear_session_context(self, key=None):
        return self.__context.clear_session_context(key)

    def init_vector_storage(self, vector_mode, collection_name="eda_service") -> Chroma:
        # embedding must use english
        embeddings = OpenAIEmbeddings()
        db_context = self.load_dataset_context()
        document_list = []
        if vector_mode == VectorMode.TABLE_SCHEMA_VECTOR:
            for table, _ in db_context["db_schema"].items():
                table_schema = db_context["db_schema"].get(table, {})
                document = "%(table_name)s"%{**table_schema}
                if table_schema.get("description"):
                    document += ": %(description)s"%{**table_schema}
                if table_schema.get("key_attributes"):
                    document += ", key_attributes = %(key_attributes)s"%{**table_schema}
                document_list.append(Document(page_content=document, metadata={"table": table}))
        else:
           for table in db_context["db_schema"].values():
               for column in table.get("columns", {}).values():
                    document = "%(name)s"%{**column}
                    if column.get("description"):
                        document += ": %(description)s"%{**column}
                    if column.get("sample_data"):
                        document += ", sample_data = %(sample_data)s"%{**column}
                    document_list.append(Document(page_content=document, metadata={"table": table.get("table_name", ""), "column": column.get("name", "")}))
                    
        return Chroma.from_documents(documents=document_list, embedding=embeddings, collection_name=collection_name)
    
context = Context()