import json
import os
import sys
import urllib.parse
from datetime import datetime
from string import Template

import numpy
import uvicorn
from fastapi import FastAPI, BackgroundTasks
from starlette.requests import Request
from starlette.responses import HTMLResponse, RedirectResponse, Response
from starlette.staticfiles import StaticFiles

from arguments import get_arguments
from web.backend import Backend
from web.templating import (
    generate_index_html,
    generate_index_html_validation,
)
from web.utils import (
    CommentBody,
    SourcingState,
    RequestBody,
    ValidationRequestBody,
    WorkingRequestBody,
)


app = FastAPI()
app.mount("/static", StaticFiles(directory="web/static"), name="static")


def get_prediction(args, example, models, processor, tokenizer):
    option_probs = []
    for model in models:
        pred_label, pred_probs = predict_example(
            args, model, example, processor, tokenizer
        )
        option_probs.append(pred_probs)
    option_average_probs = list(numpy.average(option_probs, axis=0))
    pred_label = int(numpy.argmax(option_average_probs))
    gold_label = int(example["gold_answer"][-1]) - 1
    if example["mode"] in ["adv", "plain"]:
        success = gold_label != pred_label
    else:
        success = False
    return {
        "model_prediction_probability": option_average_probs,
        "model_prediction_label": pred_label,
        "success": success,
        "gold_label": gold_label,
        "probabilities": option_probs,
    }


def make_prediction(input_body: RequestBody):
    example = input_body.convert_to_example()
    return get_prediction(
        app.state.args,
        example,
        app.state.model_kit["models"],
        app.state.model_kit["processor"],
        app.state.model_kit["tokenizer"],
    )


def call_template(template_dir, backend, uid, worker_id, run_mode, is_external):
    collection_mode = backend.user_assignments[uid].collection_mode
    template_name = backend.hit_config[collection_mode].get(
        "template_name", collection_mode
    )
    instruction_path = backend.hit_config[collection_mode].get(
        "instruction_template", f"instruction_{template_name}.html"
    )
    mini_instruction_path = backend.hit_config[collection_mode].get(
        "mini_instruction_template", f"mini_instruction_{template_name}.html"
    )
    faq_path = backend.hit_config[collection_mode].get(
        "faq_template", f"faq_{template_name}.html"
    )
    task_path = backend.hit_config[collection_mode].get(
        "task_template", f"task_{template_name}.html"
    )
    subconfig = backend.hit_config[collection_mode].get("subconfig", "")
    if collection_mode == "plain" and subconfig == "justification":
        instruction_path = "instruction_justification.html"
        faq_path = "faq_justification.html"

    index_title = backend.hit_config[collection_mode].get("index_title", "")

    # todo: migrate tutorial codes
    if collection_mode in ["plain", "adv"]:
        passages = backend.get_passages(uid)
        task_ids = backend.get_task_ids(uid)
        if not index_title:
            index_title = "Writing hard reading comprehension questions."
        return generate_index_html(
            collection_mode,
            index_title,
            os.path.join(template_dir, "index.html"),
            os.path.join(template_dir, instruction_path),
            os.path.join(template_dir, faq_path),
            os.path.join(template_dir, mini_instruction_path),
            os.path.join(template_dir, task_path),
            uid,
            task_ids,
            passages,
            backend.get_current_task_index(uid),
            backend.get_success_task_ids(uid),
            app.state.args.run_mode,
            "yes" if app.state.args.is_external else "no",
            subconfig,
        )
    elif collection_mode in ["validation"]:
        if not index_title:
            if template_name == "evaluation":
                index_title = "Evaluate reading comprehension questions!"
            else:
                index_title = "Solve reading comprehension questions!"
        validation_tasks = backend.get_user_validation_task(uid, worker_id)
        if validation_tasks is None:
            backend.logger.info(
                f"Cannot issue a new validation task for worker_id={worker_id}"
            )
            return open(os.path.join(template_dir, "error.html"), "r").read()
        return generate_index_html_validation(
            collection_mode,
            index_title,
            os.path.join(template_dir, "index.html"),
            os.path.join(template_dir, instruction_path),
            os.path.join(template_dir, faq_path),
            os.path.join(template_dir, mini_instruction_path),
            os.path.join(template_dir, task_path),
            uid,
            validation_tasks,
            app.state.args.run_mode,
            "yes" if app.state.args.is_external else "no",
            subconfig=subconfig,
        )

    else:
        raise ValueError(f"Invalid collection mode: {collection_mode}")


@app.get("/", response_class=HTMLResponse)
def read_root(
    request: Request,
    uid: str = None,
    status: str = None,
    worker_id: str = None,
    assignment_id: str = None,
    debug: str = None,
    assignmentId: str = None,
    workerId: str = None,
):
    # # emergency mode
    # return "We are sorry that no HITs are available. Please click 'Return' on the HIT page to avoid any impact on your approval rating."
    # return "We are sorry that an error has occured. Please click 'Return' on the HIT page to avoid any impact on your approval rating."

    if workerId is not None and worker_id is None:
        worker_id = workerId
    if assignmentId is not None and assignment_id is None:
        assignment_id = assignmentId

    template_dir = app.state.args.template_dir
    run_mode = app.state.args.run_mode

    backend = Backend.get_backend(app)
    # uid must be included in url when sandbox and release modes
    if not uid and run_mode in ["sandbox", "release"]:
        # backend.logger.info("Called error.html")
        return open(os.path.join(template_dir, "error.html"), "r").read()

    # if the link is clicked in preview
    if assignment_id == "ASSIGNMENT_ID_NOT_AVAILABLE" and run_mode != "debug":
        return open(os.path.join(template_dir, "not_accepted.html"), "r").read()

    datestr = datetime.now().isoformat()
    if not uid and run_mode == "debug":  # issue a uid if debug
        uid = backend.get_debug_uid()
        if debug:
            url = urllib.parse.urlencode({"uid": uid, "debug": debug})
            return RedirectResponse(url="?" + url)
        else:
            url = urllib.parse.urlencode({"uid": uid})
            return RedirectResponse(url="?" + url)

    # check how many HITs the worker did
    if app.state.args.worker_log_path and worker_id:
        if not backend.check_if_worker_can_accept_hit(worker_id, uid):
            backend.logger.info(f"GET: worker {worker_id} cannot accept hit")
            return open(os.path.join(template_dir, "known_worker.html"), "r").read()
        else:
            backend.logger.info(f"GET: worker {worker_id} can accept hit")

    # create user data
    backend.logger.info(f"GET / {datestr} / {uid} {worker_id} {assignment_id} {status}")
    uid = backend.check_user_data(uid, status, worker_id, assignment_id)
    if uid is not None and status == "new":
        redirect_url = urllib.parse.urlencode(
            {
                "uid": uid,
                "assignmentId": assignment_id,
                "workerId": worker_id,
                "status": "redirect",
            }
        )
        return RedirectResponse(url="?" + redirect_url)

    if uid == "" and run_mode != "debug":
        backend.logger.info("Called error.html")
        return open(os.path.join(template_dir, "error.html"), "r").read()

    backend.logging_current_status()
    state = backend.get_current_state(uid)
    if state == SourcingState.START or state == SourcingState.WORKING:
        return call_template(
            template_dir, backend, uid, worker_id, run_mode, app.state.args.is_external
        )
    elif state == SourcingState.FINISHED or state == SourcingState.QUIT:
        fpath = os.path.join(template_dir, "finished.html")
        template = Template(open(fpath, "r").read())
        return template.substitute(
            {"completion_code": backend.get_completion_code(uid)}
        )
    else:
        return open(os.path.join(template_dir, "error.html"), "r").read()


@app.post("/predict")
def predict(body: RequestBody):
    backend = Backend.get_backend(app)
    response = make_prediction(body)
    backend.save_submission(
        body.uid.strip(), body.passage_id, body.convert_to_example(), response
    )
    return response or {"success": True}


@app.post("/submit")
def submit(body: RequestBody):
    backend = Backend.get_backend(app)
    backend.save_submission(
        body.uid.strip(), body.passage_id, body.convert_to_example(), None
    )
    return {"success": True}


@app.post("/complete")  # complete writing HIT
async def complete(body: CommentBody, background_tasks: BackgroundTasks):
    backend = Backend.get_backend(app)
    backend.finalize_user_assignment(body.uid)
    backend.save_comment(body.uid, body.comment)
    backend.logger.info(f"Mturk submitted: {backend.get_current_state(body.uid)} {body.uid} {body.worker_id} {body.assignment_id}")
    backend.logging_current_status()
    background_tasks.add_task(
        backend.sleep_and_process_terminal_actions,
        body.uid,
        body.worker_id,
        body.assignment_id,
    )
    return {"completion_code": backend.get_completion_code(body.uid)}


@app.post("/validation_submit")  # complete validation HIT
async def validation_submit(
    body: ValidationRequestBody, background_tasks: BackgroundTasks
):
    backend = Backend.get_backend(app)
    uid = body.uid.strip()
    responses = body.responses
    mode = body.mode
    backend.save_validation_submission(body.uid.strip(), mode, responses)
    background_tasks.add_task(
        backend.sleep_and_process_terminal_actions,
        uid,
        body.worker_id,
        body.assignment_id,
    )
    backend.logging_current_status()
    backend.logger.info(f"Mturk submitted: {backend.get_current_state(body.uid)} {body.uid} {body.worker_id} {body.assignment_id}")
    return {"completion_code": backend.get_completion_code(uid)}


@app.get("/serialize")
def serialize_backend_data():
    """
    To stop the API server and resume it with the same data
    """
    backend = Backend.get_backend(app)
    backend.serialize_data()


@app.get("/show_log", include_in_schema=False)
def show_log_file(mode: str = "log", passcode: str = None):
    if app.state.args.passcode:
        backend = Backend.get_backend(app)
        res = backend.show_log(mode, passcode, app.state.args.passcode)
        return Response(content=res)
    else:
        return ""


@app.get("/set_value", include_in_schema=False)
def set_value(variable: str, value: int, passcode: str = None):
    if not app.state.args.passcode:
        return ""
    backend = Backend.get_backend(app)
    if not backend.check_passcode(passcode, app.state.args.passcode):
        return "incorrect passcode"
    if variable == "max_acceptable_hit_num":
        backend.max_acceptable_hit_num = value
        return f"set {variable}={value}"


@app.get("/get_value", include_in_schema=False)
def get_value(variable: str, passcode: str = None):
    if not app.state.args.passcode:
        return ""
    backend = Backend.get_backend(app)
    if not backend.check_passcode(passcode, app.state.args.passcode):
        return "incorrect passcode"
    if variable == "max_acceptable_hit_num":
        return f"{variable}={backend.max_acceptable_hit_num}"


@app.get("/enter_debug", include_in_schema=False)
def enter_debug(passcode: str = None):
    if not app.state.args.passcode:
        return ""
    backend = Backend.get_backend(app)
    if not backend.check_passcode(passcode, app.state.args.passcode):
        return "incorrect passcode"
    # backend.dataset["validation"]['gutenberg-Deephaven_and_Selected_Stories_&amp;_Sketches-paragraph-0152_adv_3VW6495TLOJSA6R7QJX94RMER2TYYR_0'] = backend.dataset["validation"]['gutenberg-Deephaven_and_Selected_Stories_&_Sketches-paragraph-0152_adv_3VW6495TLOJSA6R7QJX94RMER2TYYR_0']
    breakpoint()


@app.post("/working")
def update_working_status(body: WorkingRequestBody):
    uid = body.uid
    backend = Backend.get_backend(app)
    backend.update_current_time(uid)
    return {"updated": 1}


def check_config_template_paths(config, template_dir):
    def check_exists(file_name):
        filepath = os.path.join(template_dir, file_name)
        if not os.path.exists(filepath):
            print(f"Not found: {filepath}")
            return False
        return True

    for temp_type in ["instruction", "mini_instruction", "faq", "task"]:
        if f"{temp_type}_template" in config:
            assert check_exists(config[f"{temp_type}_template"])
        elif "template_name" in config:
            # TODO
            template_name = config["template_name"]
            assert check_exists(f"{temp_type}_{template_name}.html")
        else:
            assert check_exists(f"{temp_type}_{config['collection_mode']}.html")


if __name__ == "__main__":
    if sys.version_info.minor <= 6:
        print("Run with Python >= 3.7")
        exit(1)
    # setup arguments
    args = get_arguments()

    # check arguments
    if args.max_acceptable_hit_num == 1 and args.run_mode == "release":
        msg = f"Proceed with args.max_acceptable_hit_num={args.max_acceptable_hit_num}? [Y/N] "
        if input(msg) != "Y":
            raise ValueError("Aborted.")

    hit_configs = json.load(open(args.hit_config, "r"))
    setattr(args, "hit_config", {})
    for config in hit_configs:
        check_config_template_paths(config, args.template_dir)
        if "hit_template" in config:
            config["hit_template"] = os.path.join(
                args.template_dir, config["hit_template"]
            )
        args.hit_config[config["collection_mode"]] = config

    if len(args.hit_config) == 1:
        collection_mode = list(args.hit_config.keys())[0]
    else:
        collection_mode = "multiple"

    # if not args.collected_passage_path and args.run_mode in ["release", "sandbox"]:
    #     msg = f"Proceed without args.collected_passage_path? [Y/N] "
    #     if input(msg) != "Y":
    #         raise ValueError("Aborted.")

    # set up path strings
    datestr = datetime.now().strftime("%Y%m%d%H%M%S")

    if not args.worker_log_path and args.run_mode != "debug":
        args.worker_log_path = os.path.join(
            args.worker_log_dir, f"{args.run_mode}{datestr}.json"
        )

    args.log_dir = os.path.join(args.log_dir, f"{args.run_mode}", f"{collection_mode}")
    os.makedirs(args.log_dir, exist_ok=True)
    if args.resume:
        annot_datestr = os.path.basename(os.path.normpath(args.save_dir))
        log_file = f"{args.log_dir}/{annot_datestr}.log"
        if os.path.exists(log_file):
            args.log_file = log_file
        else:
            args.log_file = f"{args.log_dir}/{args.annotation}{datestr}resume.log"
    else:
        args.save_dir = os.path.join(
            args.save_dir,
            f"{args.run_mode}",
            f"{collection_mode}",
            f"{args.annotation}{datestr}",
        )
        os.makedirs(args.save_dir, exist_ok=True)
        args.log_file = f"{args.log_dir}/{args.annotation}{datestr}.log"

    if args.run_mode == "debug":
        args.is_external = True

    print(json.dumps(args.__dict__))

    for k, v in json.load(open(args.mturk_config, "r")).items():
        setattr(args, k, v)

    # setup models
    models = []
    if (
        not args.no_model
        and len(set(args.hit_config.keys()) & set(["adv"])) > 0
    ):
        for k, v in json.load(open(args.model_config, "r")).items():
            setattr(args, k, v)
        from model.run_multiple_choice import load_model
        if args.model_type == "unifiedqa":
            from model.run_multiple_choice import predict_generation_example as predict_example
        else:
            from model.run_multiple_choice import predict_example

        if args.model_ensemble:
            for model_path in args.ensemble_models:
                args.model_name_or_path = model_path
                model, processor, tokenizer = load_model(args)
                models.append(model)
        else:
            model, processor, tokenizer = load_model(args)
            models.append(model)
    else:
        processor, tokenizer = None, None

    app.state.model_kit = {
        "models": models,
        "processor": processor,
        "tokenizer": tokenizer,
    }

    app.state.args = args

    if args.is_external:
        question_template_path = None
    else:
        question_template_path = os.path.join(
            args.template_dir, "hit_template_COLLECTION_MODE.html"
        )
    if not args.base_url.endswith(":8000"):
        if args.base_url.endswith("/"):
            args.base_url = args.base_url[:-1]
        args.base_url = args.base_url + f":{args.port}"
    mturk_config = {
        "access_id": args.aws_access_key_id,
        "secret_key": args.aws_secret_access_key,
        "base_url": args.base_url,
        "template_path": question_template_path,
        "issued_hits": "issued_hits",  # save directory of issued HITs
    }
    app.state.mturk_config = mturk_config

    Backend.get_backend(app)

    # start app
    if args.run_mode in ["sandbox", "release"] and "localhost" not in args.base_url:
        uvicorn.run(app, host="0.0.0.0", port=args.port, log_level="info")
    else:
        uvicorn.run(app, log_level="info")
