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 openai.types.chat import ChatCompletionMessage
from pydantic import BaseModel

OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")


class ZipCode(BaseModel):
    value: int = 30328


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):
    _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.
    """

    zip: ZipCode
    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: ChatCompletionMessage) -> 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() if key != "zip"
    ]
    query_string = "&".join(query_parameters)
    zip_code = arguments.get("zip", {"value": 30328})

    endpoint = "https://www.mbusa.com/en/inventory/search"
    url = f"{endpoint}?zip={zip_code['value']}&{query_string}"

    return url


def main(openai_api_key: str, prompt: str) -> str:
    """Returns a valid URL for the Mercedes-Benz USA website based on the user's query."""

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

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

    # Connect to OpenAI with the provided API key
    client = OpenAI(api_key=openai_api_key)
    # Define an OpenAI tool for guaranteeing the schema of the response
    tool = [{"type": "function", "function": CarParameters.openai_schema}]

    # Call the OpenAI API
    response = client.chat.completions.create(
        model="gpt-4-turbo",
        messages=messages,
        temperature=0,
        tools=tool,  # Ensure the response schema is correct
        tool_choice="required",  # Ensure the tool is used
    )

    # Parse the response and construct the URL
    url = format_url(response)
    markdown_link = f"[{url}]({url})"  # Return clickable link for convenience
    return url, markdown_link


with gr.Blocks() as demo:
    with gr.Column():
        with gr.Row():
            openai_api_key = gr.Textbox(
                label="api-key",
                placeholder="Set your OpenAI API key here",
                lines=1,
                type="password",
            )
        with gr.Row():
            with gr.Column():
                prompt = gr.Textbox(
                    label="prompt",
                    placeholder="Enter your query here and press Enter",
                    lines=1,
                )
            with gr.Column():
                url = gr.Textbox(
                    label="url",
                    placeholder="URL will appear here",
                    lines=1,
                    type="text",
                )
        with gr.Row():
            link = gr.Markdown()
        with gr.Row():
            submit = gr.Button(value="Submit")

    submit.click(fn=main, inputs=[openai_api_key, prompt], outputs=[url, link])

    examples = [
        "I need a car for my growing family",
        "I frequently travel in snowy conditions",
        "I'm an eco-conscious city dweller",
        "I love high-performance sports cars",
        "I live in the snowy tundra's of Alaska",
        "Im looking for a 2024 Mercedes-Benz AMG SUV with a black exterior and interior, and 17 inch wheels. Can you help me with that?",
    ]

    gr.Examples(examples=examples, inputs=prompt, outputs=url)
demo.launch(inline=True, debug=True)