"""
    :author: Jinfen Li
    :url: https://github.com/JinfenLi
"""
import socket
import subprocess
from collections import defaultdict as ddict
import requests
from tqdm import tqdm

def remap_token_id(reference_offset, parsed_offset, ref_offset_tid = 1):
    """

    Args:
        reference_offset: offset generated by original model
        parsed_offset: offset generated by stanfordnlp
        ref_offset_tid: initial reference_offset token id, the first reference_offset is for [CLS]

    Returns: {new_token_dict: [start_ref_tid, end_ref_tid)}, end_ref_tid is not inclusive

    """
    parsed_offset_tid = 0
    new_token_dict = {}
    start = ref_offset_tid
    while parsed_offset_tid < len(parsed_offset):
        if reference_offset[ref_offset_tid][0] == 0 and reference_offset[ref_offset_tid][1] == 0:
            break
        elif parsed_offset[parsed_offset_tid][1] >= reference_offset[ref_offset_tid][1]:
            # find the first reference_offset that exceed parsed_offset
            while ref_offset_tid < (len(reference_offset) - 1) and parsed_offset[parsed_offset_tid][1] > reference_offset[ref_offset_tid][
                1]:
                ref_offset_tid += 1
            if parsed_offset_tid > 0:
                start = new_token_dict[parsed_offset_tid - 1][1]
            end = ref_offset_tid + 1
            new_token_dict[parsed_offset_tid] = [start, end]
            parsed_offset_tid += 1
        elif parsed_offset[parsed_offset_tid][1] < reference_offset[ref_offset_tid][1]:
            if parsed_offset_tid > 0:
                start = new_token_dict[parsed_offset_tid - 1][1]
            end = ref_offset_tid + 1
            new_token_dict[parsed_offset_tid] = [start, end]
            parsed_offset_tid += 1

    return new_token_dict, ref_offset_tid + 1

def parse_senti_tree(senti_tree_str, tid=0):
    """
    Args:
        senti_tree_str: a string of sentiment tree
        tid: initial token id
    Returns: phrase_dict: {'sentiment': [1, 2], 'tokens': [[1, 2], [1, 2]]}

    """
    senti_nodes = [chars.strip() for chars in senti_tree_str.split('\r\n')]
    phrase_dict = ddict(list)
    # stack: [{label: str, sentiment: int, span: list, left_count: int, right_count: int}]

    stack = []

    for node in senti_nodes:
        labels = [l.split("|")[0] for l in node.split('(') if l]
        sentiments = [int(s) for s in re.findall(r"sentiment=(\d)", node)]

        if len(labels) == 3:
            stack.append({"label": labels[0], "sentiment": sentiments[0], "span": [tid, tid + 1], "left_count": node.count("("), "right_count": node.count(")")})
            tid += 2
        elif len(labels) == 2:
            stack.append(
                {"label": labels[0], "sentiment": sentiments[0], "span": [tid], "left_count": node.count("("),
                 "right_count": node.count(")")})
            tid += 1
        elif len(labels) == 1:
            stack.append({"label": labels[0], "sentiment": sentiments[0], "span": [tid] if node.count(")") > 0 else None, "left_count": node.count("("), "right_count": node.count(")")})
            if node.count(")") > 0:
                tid += 1
        # if right_count > left_count, means rolling back
        while stack and stack[-1]["right_count"] > stack[-1]["left_count"]:
            element = stack.pop()
            stack[-1]["right_count"] = stack[-1]["right_count"] + element["right_count"]
            stack[-1]["left_count"] = stack[-1]["left_count"] + element["left_count"]
            # the last element is the left node
            if stack[-1]["span"] is not None and len(stack[-1]["span"]) == 2:
                continue
            stack[-1]["span"] = [stack[-1]["span"][0] if stack[-1]["span"] is not None else element["span"][0], tid - 1]
            if stack[-1]["label"] in ['@S', 'S', 'ROOT']:
                phrase_dict['sentiment'].append(stack[-1]["sentiment"])
                phrase_dict['tokens'].append(stack[-1]["span"])

    return phrase_dict

def sentiment_tree(texts, num_examples, original_offsets, resource_dir):
    sents_with_tokens_list = []
    available_port = get_CoreNLPClient(resource_dir)
    # with CoreNLPClient(be_quiet=False, annotators=['tokenize', 'sentiment'], properties={'tokenize.codepoint': 'true'}, timeout=60000,
    #                    memory='16G', output_format='json') as client:

    for idx in tqdm(range(num_examples), desc=f"Building dataset"):
        # new_senti_dict = ddict(list)
        text = texts[idx]
        original_offset = original_offsets[idx]
        # ann = client.annotate(text.strip())
        ann = annotate_text(text.strip(), available_port)
        cur_start_tid = 0
        result_dict = ddict(list)
        new_token_dict, phrase_dict = {}, {}
        cur_start_ref_tid = 1
        for sent_dict in ann['sentences']:

            sent_dict['offset'] = [(x['characterOffsetBegin'], x['characterOffsetEnd']) for x in sent_dict['tokens']]
            # original_token_idx = [x['index'] for x in sent_dict['tokens']]
            remap_token_dict, cur_start_ref_tid = remap_token_id(original_offset, sent_dict['offset'], cur_start_ref_tid)
            new_token_dict.update({k + cur_start_tid: v for k, v in remap_token_dict.items()})
            phrase_dict = parse_senti_tree(sent_dict['sentimentTree'], cur_start_tid)
            phrase_dict['tokens'] = [[new_token_dict[x[0]][0], new_token_dict[x[1]][1]] for x in phrase_dict['tokens']]
            result_dict['spans'].extend(phrase_dict['tokens'])
            result_dict['sentiment'].extend(phrase_dict['sentiment'])
            cur_start_tid = max(new_token_dict.keys()) + 1
        sents_with_tokens_list.append(result_dict)


    return sents_with_tokens_list


def get_CoreNLPClient(resource_dir):
    def find_available_port(start_port, max_attempts=10):
        for port in range(start_port, start_port + max_attempts):
            try:
                with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
                    s.bind(("127.0.0.1", port))
                return port
            except OSError:
                pass
        return None

    available_port = find_available_port(start_port=9000)

    if available_port:
        print("Available port for CoreNLPClient:", available_port)
    else:
        print("No available port  for CoreNLPClient found")
    # Define the command to start the CoreNLP server
    corenlp_command = [
        "java", "-mx16g", "-cp", f"{resource_dir}/stanford-corenlp-4.5.4/*",
        "edu.stanford.nlp.pipeline.StanfordCoreNLPServer",
        "-port", f"{available_port}", "-timeout", "60000", "-threads", "4", "-maxCharLength", "100000", "-quiet",
        "True",
        "-annotators", "tokenize,sentiment",
        "properties", "{'tokenize.codepoint': 'true'}"
    ]

    # Start the CoreNLP server
    subprocess.Popen(corenlp_command)
    return available_port

def annotate_text(text, available_port):

    corenlp_url = f"http://localhost:{available_port}"
    properties = {
        "annotators": "parse,sentiment",
        "properties": "{'tokenize.codepoint': 'true'}",
        "outputFormat": "json"
    }

    response = requests.post(f"{corenlp_url}/?properties={properties}", data=text.encode("utf-8"))
    json_response = response.json()
    return json_response
