import os
import traceback

from pymongo import ReturnDocument

from ...pipelines.base import BasePipeline
from ...connectors.databases.mongodb_connector import MongoDBConnector
from ...connectors.message_brokers.rabbitmq_connector import RabbitMQConnector

class WebsitePipeline(BasePipeline):

    def __init__(self):

        self.receive_queue = None
        self.transmit_queue = None

        self.website_coll = None
        self.categories_coll = None

        super().__init__()

    def load_database_env(self):
        super().load_database_env()


    def setup_database(self, reinit=False):
        super().setup_database()
        self.website_coll = self.database[f'{os.getenv("WEBSITE_COLL")}']
        self.categories_coll = self.database[f'{os.getenv("CATEGORIES_COLL")}']

    def load_connector_env(self):
        super().load_connector_env()
        self.receive_queue = os.getenv("RECEIVE_QUEUE")
        self.connector_params['receive_queue'] = self.receive_queue
        self.transmit_queue = os.getenv("TRANSMIT_QUEUE")
        self.connector_params['transmit_queue'] = self.transmit_queue

    def init_database_connector(self):
        self.database_client = MongoDBConnector(
            connection_string=self.database_params['conn_string'],
            default_database=self.database_params['database']
        )
        self.database = self.database_client.connection[self.database_params['database']]

    def init_message_connector(self):
        self.connector_client = RabbitMQConnector()

    def process_item(self, item, spider):
        # update status of website_data
        try:
            website = None
            with self.database_client.connection.start_session() as session:
                with session.start_transaction():
                    update_obj = None

                    if item["status"] != "error":
                        update_obj = {
                            "name": item["name"],
                            "brand": item["brand"],
                            "status": item["status"],
                            "last_crawl": item["last_crawl"],
                            "css_selector": item["css_selector"],
                            "categories_num": item["categories_num"],
                            "total_categories": item["total_categories"],
                            "err_re_crawl": 0,
                            "err_re_process": 0,  # clear err when cron init
                        }
                    else:
                        website = self.website_coll.find_one({"url": item["url"]})

                        update_obj = {
                            "name": item["name"],
                            "brand": item["brand"],
                            "status": item["status"],
                            "last_crawl": item["last_crawl"],
                            "total_categories": item["total_categories"],
                            "err_re_crawl": website["err_re_crawl"] + 1,
                            "err_re_process": 0,  # clear err when cron init
                        }

                    website = self.website_coll.find_one_and_update(
                        {"url": item["url"]},
                        {
                            "$set": update_obj,
                        },
                        return_document=ReturnDocument.AFTER,
                        session=session,
                    )

            # send all categories_url to categories_queue

            if website:
                for category in item["category_urls"]:
                    # test rabbit
                    if category != item["category_urls"][0]:
                        continue

                    body_message = {
                        "website_id": str(website["_id"]),
                        "website_url": website["url"],
                        "website_name": website["name"],
                        "url": category,
                        "spider_name": website["spider_name"],
                    }

                    # send message to rabbitmq (categories_queue)
                    self.send_message(body_message, self.connector_params['transmit_queue'])

            self.get_logger().info(
                f'[*] Crawling website {item["url"]} completed at {item["last_crawl"]}: status - "{item["status"]}"'
            )
        except Exception as error:
            # traceback.print_exc()
            # + log to log server
            self.get_logger().error(
                f'[x] Error when process data crawled in pipeline website (url: {item["url"]}). \n'
                + traceback.format_exc()
            )

        return item

    def close_spider(self, spider):
        # disconnect database
        self.database_client.close()

        # disconnect rabbitmq
        self.connector_client.close()