from fastapi import FastAPI, status, Response
from agent import agent, system_prompt
from functions import process_input_json, evalute_final_json, get_variables_and_boundaries
import traceback


app = FastAPI()


@app.post("/generate-step")
async def generate_step(input_json: dict, response: Response = None):
    try:
        rl_input_json = input_json['rl_input_json']
        usecase_details = input_json['step_json']

        if "title" not in usecase_details or not usecase_details['title']:
            response.status_code = status.HTTP_422_UNPROCESSABLE_ENTITY
            return {"code": 422, "data": "Usecase details is required.", "success": False}
        if "observableFactor" not in rl_input_json or not rl_input_json['observableFactor']:
            response.status_code = status.HTTP_422_UNPROCESSABLE_ENTITY
            return {"code": 422, "data": "Observable factor is required.", "success": False}
        if "actions" not in rl_input_json or not rl_input_json['actions']:
            response.status_code = status.HTTP_422_UNPROCESSABLE_ENTITY
            return {"code": 422, "data": "Actions is required.", "success": False}
        if "boundaries" not in rl_input_json or not rl_input_json['boundaries']:
            response.status_code = status.HTTP_422_UNPROCESSABLE_ENTITY
            return {"code": 422, "data": "Boundaries is required.", "success": False}

        processed_json = process_input_json(rl_input_json['observableFactor'], rl_input_json.get(
            'constantFactor'), rl_input_json.get('modelSelection'), rl_input_json['actions'], rl_input_json['boundaries'])

        print("rl_step | processed_json:", processed_json)
        try:
            evalute_final_json(processed_json)
        except Exception:
            trace = traceback.format_exc()
            if "KeyError:" in trace:
                response.status_code = status.HTTP_422_UNPROCESSABLE_ENTITY
                return {"code": 422, "data": trace, "success": False}

        details = get_variables_and_boundaries(processed_json)

        last_trace = None
        current_errors = ""

        current_prompt = system_prompt
        existing_reward = ""
        if "reward" in rl_input_json and rl_input_json['reward']["selectedSection"] == "ai" and rl_input_json['reward']['generatedCode']:
            current_prompt += "\n###Existing reward:\n```" + \
                rl_input_json['reward']['generatedCode'].replace(
                    "{", "{{").replace("}", "}}") + "```"
            existing_reward = rl_input_json['reward']['generatedCode']
        elif "reward" in rl_input_json and rl_input_json['reward']["selectedSection"] != "ai" and rl_input_json['reward']['customCode']:
            current_prompt += "\n###Existing reward:\n```" + \
                rl_input_json['reward']['customCode'].replace(
                    "{", "{{").replace("}", "}}") + "```"
            existing_reward = rl_input_json['reward']['customCode']

        for _ in range(3):
            final_prompt = current_prompt + "\n" + current_errors

            try:
                agent_response = agent.invoke({"input": final_prompt.format(
                    json_data=processed_json, use_case_name=usecase_details['title'], variable_information=details, use_case_description=usecase_details['description'])})

                step = agent_response['output'].get("step")
                reward = agent_response['output'].get("reward")

                processed_json['environment']['step'] = ["\n".join(step)]
                processed_json['environment']['reward'] = [
                    "\n".join(reward)] if not existing_reward else [existing_reward]

                try:
                    print(
                        "##################### STARTED EVALUATION ##############################")
                    evalute_final_json(processed_json)
                except Exception:
                    print(
                        "######################## Error in evalution #########################################")
                    last_trace = traceback.format_exc()

                    current_errors = "In last executation I got below error so make sure you I don't get same error again. \nError: \n" + last_trace
                    print(last_trace)
                    print(
                        "######################## Error in evalution #########################################")
                    continue
                response.status_code = status.HTTP_200_OK
                return {"code": 200, "data": "\n".join(agent_response['output']['step']), "success": True}

            except Exception:
                last_trace = traceback.format_exc()
                print("rl_step | Error trace:", last_trace)
        response.status_code = status.HTTP_500_INTERNAL_SERVER_ERROR
        return {"code": 500, "data": last_trace, "success": False}

    except Exception:
        trace = traceback.format_exc()
        print("rl_step | Error trace:", trace)
        if "KeyError:" in trace:
            response.status_code = status.HTTP_422_UNPROCESSABLE_ENTITY
            return {"code": 422, "data": trace, "success": False}


@app.post("/generate-reward")
async def generate_reward(input_json: dict, response: Response = None):
    try:
        rl_input_json = input_json['rl_input_json']
        usecase_details = input_json['reward_json']

        if "title" not in usecase_details or not usecase_details['title']:
            response.status_code = status.HTTP_422_UNPROCESSABLE_ENTITY
            return {"code": 422, "data": "Usecase details is required.", "success": False}
        if "observableFactor" not in rl_input_json or not rl_input_json['observableFactor']:
            response.status_code = status.HTTP_422_UNPROCESSABLE_ENTITY
            return {"code": 422, "data": "Observable factor is required.", "success": False}
        if "actions" not in rl_input_json or not rl_input_json['actions']:
            response.status_code = status.HTTP_422_UNPROCESSABLE_ENTITY
            return {"code": 422, "data": "Actions is required.", "success": False}
        if "boundaries" not in rl_input_json or not rl_input_json['boundaries']:
            response.status_code = status.HTTP_422_UNPROCESSABLE_ENTITY
            return {"code": 422, "data": "Boundaries is required.", "success": False}

        processed_json = process_input_json(rl_input_json['observableFactor'], rl_input_json.get(
            'constantFactor'), rl_input_json.get('modelSelection'), rl_input_json['actions'], rl_input_json['boundaries'])

        print("rl_reward | processed_json:", processed_json)
        try:
            evalute_final_json(processed_json)
        except Exception:
            trace = traceback.format_exc()
            print("rl_reward | Error trace:", trace)
            if "KeyError:" in trace:
                response.status_code = status.HTTP_422_UNPROCESSABLE_ENTITY
                return {"code": 422, "data": trace, "success": False}

        details = get_variables_and_boundaries(processed_json)

        last_trace = None
        current_errors = ""

        current_prompt = system_prompt

        existing_step = ""
        if "step" in rl_input_json and rl_input_json['step']["selectedSection"] == "ai" and rl_input_json['step']['generatedCode']:
            current_prompt += "\n###Existing step:\n```" + \
                rl_input_json['step']['generatedCode'].replace(
                    "{", "{{").replace("}", "}}") + "```"
            existing_step = rl_input_json['step']['generatedCode']
        elif "step" in rl_input_json and rl_input_json['step']["selectedSection"] != "ai" and rl_input_json['step']['customCode']:
            current_prompt += "\n###Existing step:\n```" + \
                rl_input_json['step']['customCode'].replace(
                    "{", "{{").replace("}", "}}") + "```"
            existing_step = rl_input_json['step']['customCode']

        for _ in range(3):
            final_prompt = current_prompt + "\n" + current_errors

            try:
                agent_response = agent.invoke({"input": final_prompt.format(
                    json_data=processed_json, use_case_name=usecase_details['title'], variable_information=details, use_case_description=usecase_details['description'])})

                step = agent_response['output'].get("step")
                reward = agent_response['output'].get("reward")

                processed_json['environment']['step'] = [
                    "\n".join(step)] if not existing_step else [existing_step]
                processed_json['environment']['reward'] = ["\n".join(reward)]

                try:
                    print(
                        "##################### STARTED EVALUATION ##############################")
                    evalute_final_json(processed_json)
                except Exception:
                    print(
                        "######################## Error in evalution #########################################")
                    last_trace = traceback.format_exc()

                    current_errors = "In last executation I got below error so make sure you I don't get same error again. \nError: \n" + last_trace
                    print(last_trace)
                    print(
                        "######################## Error in evalution #########################################")
                    continue
                response.status_code = status.HTTP_200_OK
                return {"code": 200, "data": "\n".join(agent_response['output']['reward']), "success": True}

            except Exception:
                last_trace = traceback.format_exc()
                print("rl_reward | Error trace:", last_trace)
        response.status_code = status.HTTP_500_INTERNAL_SERVER_ERROR
        return {"code": 500, "data": last_trace, "success": False}

    except Exception:
        trace = traceback.format_exc()
        print("rl_reward | Error trace:", trace)
        if "KeyError:" in trace:
            response.status_code = status.HTTP_422_UNPROCESSABLE_ENTITY
            return {"code": 422, "data": trace, "success": False}