import json
import os
from datetime import datetime
from enum import Enum
from typing import List

import gradio as gr
from instructor import OpenAISchema
from openai import OpenAI
from pydantic import BaseModel

OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")


class Bodystyles(Enum):
    CAB = "CAB"
    CPE = "CPE"
    RDS = "RDS"
    SDN = "SDN"
    SUV = "SUV"
    WGN = "WGN"


class Brands(Enum):
    AMG = "AMG"
    MAYBACH = "MAYBACH"


class ColorsExterior(Enum):
    BLK = "BLK"
    BLU = "BLU"
    BWN = "BWN"
    GRN = "GRN"
    GRY = "GRY"
    OTR = "OTR"
    RED = "RED"
    SLV = "SLV"
    WHT = "WHT"
    YLW = "YLW"


class ColorsInterior(Enum):
    BGE = "BGE"
    BLK = "BLK"
    BWN = "BWN"
    GRY = "GRY"
    OTR = "OTR"
    RED = "RED"


class DistanceFilter(Enum):
    _10 = "10"
    _100 = "100"
    _1000 = "1000"
    _200 = "200"
    _25 = "25"
    _50 = "50"
    _500 = "500"
    ANY = "ANY"


class Engines(Enum):
    # Adding only a few options for brevity
    _20LINLINE4TURBO = "20LINLINE4TURBO"
    _20LINLINE4TURBOWITHMILDHYBRIDDRIVE = "20LINLINE4TURBOWITHMILDHYBRIDDRIVE"
    _20LINLINE4TURBOWITHPLUGINHYBRIDTECHNOLOGY = (
        "20LINLINE4TURBOWITHPLUGINHYBRIDTECHNOLOGY"
    )
    _30LINLINE6TURBOENGINEWITHHYBRIDASSIST = "30LINLINE6TURBOENGINEWITHHYBRIDASSIST"
    _30LINLINE6TURBOENGINEWITHMILDHYBRIDDRIVE = (
        "30LINLINE6TURBOENGINEWITHMILDHYBRIDDRIVE"
    )
    _30LINLINE6TURBOENGINEWITHMILDHYBRIDDRIVEANDELECTRICAUXILIARYCOMPRESSOR = (
        "30LINLINE6TURBOENGINEWITHMILDHYBRIDDRIVEANDELECTRICAUXILIARYCOMPRESSOR"
    )
    _30LINLINE6TURBOENGINEWITHMILDHYBRIDDRIVEPLUSPLUGINHYBRIDELECTRICMOTOR = (
        "30LINLINE6TURBOENGINEWITHMILDHYBRIDDRIVEPLUSPLUGINHYBRIDELECTRICMOTOR"
    )
    _30LINLINE6TURBOWITHHYBRIDASSIST = "30LINLINE6TURBOWITHHYBRIDASSIST"
    _30LINLINE6TURBOWITHMILDHYBRIDDRIVE = "30LINLINE6TURBOWITHMILDHYBRIDDRIVE"
    _40LV8BITURBO = "40LV8BITURBO"
    _40LV8BITURBOWITHMILDHYBRIDDRIVE = "40LV8BITURBOWITHMILDHYBRIDDRIVE"
    _AMGENHANCED20LINLINE4TURBO = "AMGENHANCED20LINLINE4TURBO"
    _AMGENHANCED20LINLINE4TURBOWITHHYBRIDASSIST = (
        "AMGENHANCED20LINLINE4TURBOWITHHYBRIDASSIST"
    )
    _AMGENHANCED30LINLINE6TURBOWITHHYBRIDASSIST = (
        "AMGENHANCED30LINLINE6TURBOWITHHYBRIDASSIST"
    )
    _AMGENHANCED30LINLINE6TURBOWITHHYBRIDASSISTANDELECTRICAUXILIARYCOMPRESSOR = (
        "AMGENHANCED30LINLINE6TURBOWITHHYBRIDASSISTANDELECTRICAUXILIARYCOMPRESSOR"
    )
    _AMGENHANCED30LV6BITURBO = "AMGENHANCED30LV6BITURBO"
    _DUALPERMANENTLYEXCITEDSYNCHRONOUSPSM = "DUALPERMANENTLYEXCITEDSYNCHRONOUSPSM"
    _FRONTASYNCHRONOUSASMREARPERMANENTLYEXCITEDSYNCHRONOUSPSM = (
        "FRONTASYNCHRONOUSASMREARPERMANENTLYEXCITEDSYNCHRONOUSPSM"
    )
    _FRONTMOUNTEDASYNCHRONOUSASM = "FRONTMOUNTEDASYNCHRONOUSASM"
    _HANDCRAFTED20LINLINE4TURBOWITHHYBRIDASSIST = (
        "HANDCRAFTED20LINLINE4TURBOWITHHYBRIDASSIST"
    )
    _HANDCRAFTED60LV12BITURBO = "HANDCRAFTED60LV12BITURBO"
    _HANDCRAFTEDAMG20LINLINE4TURBO = "HANDCRAFTEDAMG20LINLINE4TURBO"
    _HANDCRAFTEDAMG20LINLINE4TURBOWITHHYBRIDASSISTANDELECTRICEXHAUSTGASTURBOCHARGER = (
        "HANDCRAFTEDAMG20LINLINE4TURBOWITHHYBRIDASSISTANDELECTRICEXHAUSTGASTURBOCHARGER"
    )
    _HANDCRAFTEDAMG40LV8BITURBO = "HANDCRAFTEDAMG40LV8BITURBO"
    _HANDCRAFTEDAMG40LV8BITURBOWITHHYBRIDASSIST = (
        "HANDCRAFTEDAMG40LV8BITURBOWITHHYBRIDASSIST"
    )
    _PERMANENTLYEXCITEDSYNCHRONOUSPSM = "PERMANENTLYEXCITEDSYNCHRONOUSPSM"


class FeaturesOptions(Enum):
    _X_OPT_810 = "X_OPT_810"
    _X_OPT_811 = "X_OPT_811"
    _X_OPT_APA = "X_OPT_APA"
    _X_OPT_APE = "X_OPT_APE"
    _X_OPT_CAM = "X_OPT_CAM"
    _X_OPT_CFC = "X_OPT_CFC"
    _X_OPT_DIC = "X_OPT_DIC"
    _X_OPT_HSW = "X_OPT_HSW"
    _X_OPT_HTS = "X_OPT_HTS"
    _X_OPT_HVS = "X_OPT_HVS"
    _X_OPT_IDS = "X_OPT_IDS"
    _X_OPT_LSW = "X_OPT_LSW"
    _X_OPT_MCS = "X_OPT_MCS"
    _X_OPT_MSC = "X_OPT_MSC"
    _X_OPT_PAR = "X_OPT_PAR"
    _X_OPT_RSS = "X_OPT_RSS"
    _X_OPT_RWS = "X_OPT_RWS"
    _X_OPT_SBC = "X_OPT_SBC"
    _X_OPT_WLSW = "X_OPT_WLSW"
    _X_OPT_WPC = "X_OPT_WPC"


class FeaturesPackages(Enum):
    _X_PKG_318 = "X_PKG_318"
    _X_PKG_ALE = "X_PKG_ALE"
    _X_PKG_ANP = "X_PKG_ANP"
    _X_PKG_AP = "X_PKG_AP"
    _X_PKG_APS = "X_PKG_APS"
    _X_PKG_DIP = "X_PKG_DIP"
    _X_PKG_E321 = "X_PKG_E321"
    _X_PKG_EIP = "X_PKG_EIP"
    _X_PKG_PAP = "X_PKG_PAP"
    _X_PKG_PMP = "X_PKG_PMP"
    _X_PKG_S325 = "X_PKG_S325"


class FeaturesWheels(Enum):
    _X_WHL_17i = "X_WHL_17i"
    _X_WHL_18i = "X_WHL_18i"
    _X_WHL_192i = "X_WHL_192i"
    _X_WHL_19i = "X_WHL_19i"
    _X_WHL_20i = "X_WHL_20i"
    _X_WHL_21i = "X_WHL_21i"
    _X_WHL_22i = "X_WHL_22i"
    _X_WHL_23i = "X_WHL_23i"


class FuelType(Enum):
    _D = "D"
    _E = "E"
    _G = "G"
    _H = "H"


class HighwayFuelEconomy(Enum):
    _1 = "1"
    _2 = "2"
    _3 = "3"
    _4 = "4"


class PassengerCapacity(Enum):
    _4 = "4"
    _5 = "5"
    _7 = "7"


class PriceRanges(Enum):
    _1 = "50000_70000"
    _2 = "70000_90000"
    _3 = "90000_120000"
    _4 = "120000_999000"
    _5 = "0_50000"


class Years(Enum):
    _2023 = "2023"
    _2024 = "2024"


class CarParameters(OpenAISchema):
    """
    defines various query parameters used
    to search vehicle inventory on https://mbusa.com/en/inventory/search.
    """

    bodyStyle: List[Bodystyles] = []
    brand: List[Brands] = []
    exteriorColor: List[ColorsExterior] = []
    interiorColor: List[ColorsInterior] = []
    distance: List[DistanceFilter] = []
    engine: List[Engines] = []
    interiorFeature: List[FeaturesOptions] = []
    performanceSafetyFeature: List[FeaturesOptions] = []
    wheel: List[FeaturesWheels] = []
    fuelType: List[FuelType] = []
    highwayFuelEconomy: List[HighwayFuelEconomy] = []
    passengerCapacity: List[PassengerCapacity] = []
    price: List[PriceRanges] = []
    year: List[Years] = []


def format_url(response) -> str:
    tool_call = response.choices[0].message.tool_calls[0]

    arguments_string = tool_call.function.arguments
    arguments = json.loads(arguments_string)

    query_parameters = [f'{key}={",".join(value)}' for key, value in arguments.items()]
    query = "&".join(query_parameters)

    endpoint = "https://www.mbusa.com/en/inventory/search"
    url = f"{endpoint}?zip=30328&{query}"

    return url


def main(prompt: str):
    client = OpenAI(api_key=OPENAI_API_KEY)
    tool = [{"type": "function", "function": CarParameters.openai_schema}]

    system_prompt = f"""
        You are an AI assistant who specializes in generating relevant query parameters for the Mercedes-Benz USA website.
        current date: {datetime.now().strftime("%Y-%m-%d")}
    """.strip()

    user_prompt = "Im looking for a 2022 Mercedes-Benz AMG SUV \
        with a black exterior and interior, \
        and 17 inch wheels. \
        Can you help me with that?"

    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": user_prompt},
    ]

    response = client.chat.completions.create(
        model="gpt-4-turbo",
        messages=messages,
        temperature=0,
        tools=tool,
        tool_choice="required",
    )
    return format_url(response)


interface = gr.Interface(fn=main, inputs="text", outputs="text")
interface.launch()