import json

import requests


class AMBITResult:

    def __init__(self, *args, **kwargs):
        self.smiles = kwargs['smiles']
        self.tps = []
        for bt in kwargs['products']:
            if len(bt['products']):
                self.tps.append(bt)

        self.probs = None

    def __str__(self):
        x = self.smiles + "\n"
        total_bts = len(self.tps)

        for i, tp in enumerate(self.tps):
            prob = ""
            if self.probs:
                prob = f" (p={self.probs[tp['id']]})"

            if i == total_bts - 1:
                x += f"\t└── {tp['name']}{prob}\n"
            else:
                x += f"\t├── {tp['name']}{prob}\n"

            total_products = len(tp['products'])
            for j, p in enumerate(tp['products']):
                if j == total_products - 1:
                    if i == total_bts - 1:
                        x += f"\t\t└── {p}"
                    else:
                        x += f"\t│\t└── {p}\n"
                else:
                    if i == total_bts - 1:
                        x += f"\t\t├── {p}\n"
                    else:
                        x += f"\t│\t├── {p}\n"
        return x

    def set_probs(self, probs):
        self.probs = probs


class AMBIT:

    def __init__(self, host, rules=None):
        self.host = host
        self.rules = rules
        self.ambit_params = {
            'singlePos': True,
            'split': False,
        }

    def batch_apply(self, smiles: list):
        payload = {
            'smiles': smiles,
            'rules': self.rules,
        }
        payload.update(**self.ambit_params)

        res = self._execute(payload)

        tps = list()
        for r in res['result']:
            ar = AMBITResult(**r)
            if len(ar.tps):
                tps.append(ar)
            else:
                tps.append(None)
        return tps

    def apply(self, smiles: str):
        return self.batch_apply([smiles])[0]

    def _execute(self, payload):
        res = requests.post(self.host + '/ambit', data=json.dumps(payload))
        res.raise_for_status()
        return res.json()