import json
import os
from typing import Dict

import gradio as gr
from openai import OpenAI
from opensearchpy import OpenSearch

host = "localhost"
port = 9200
OPENSEARCH_ADMIN_PASSWORD = os.getenv("OPENSEARCH_ADMIN_PASSWORD", "yw7L5u9nLs3a")
auth = (
    "admin",
    OPENSEARCH_ADMIN_PASSWORD,
)  # For testing only. Don't store credentials in code.

# Create the client with SSL/TLS enabled, but hostname verification disabled.
os_client = OpenSearch(
    hosts=[{"host": host, "port": port}],
    http_compress=True,  # enables gzip compression for request bodies
    http_auth=auth,
    use_ssl=True,
    verify_certs=False,
    ssl_assert_hostname=False,
    ssl_show_warn=False,
)

all_props = json.load(open("all_properties.json"))
props_core = [
    {
        "property": prop["property"],
        "address": prop["address"],
        "school": prop["school"],
        "listPrice": prop["listPrice"],
    }
    for prop in all_props
]

client = OpenAI(
    # This is the default and can be omitted
    api_key=os.environ.get("OPENAI_API_KEY"),
)

REQUIREMENTS_KEYS = [
    "location",
    "budget",
    "house type",
    "layout",
]

requirements: Dict[str, str] = {}


def get_requirement_prompt():
    return f"Current collected requirements: {requirements}.\nPlease let me know your requirements: {[key for key in REQUIREMENTS_KEYS if key not in requirements.keys()]}"


def get_requirements(input_str):
    global requirements
    chat_completion = client.chat.completions.create(
        messages=[
            {
                "role": "system",
                "content": """
You are a real estate agent looking for properties for your clients. You are now extracting information to understand user's requests.
Requirements are:

* location. For example:
- City: Mountainview, San Jose, Dublin, etc
- Postal Code: 95134, etc
- Work or Regular Destination: work at Google; regular business flight
- School: a specific school name; specify a rating range for schools.

* budget. For example:
- around 1 million; no more than 1.5 million; between 800k to 1 million

* layout. For example:
- Bedroom: 3 bedrooms; no less than 2 bedrooms; single; married; 2 kids
- Bathroom: same as above

* house type (one or multiple choices). For example:
- condo, townhouse, single family

If user's input is a requirement, please provide the content in the format '{requirement_type}:{content}'. For example, 'location:Houston'. If multiple requirements are provided, separate them with |. For example, 'location:Houston|budget:1 million'.
If multiple requirements for the same type are provided, separate them with ;. For example, 'location:Houston;San Francisco'. DO NOT output multiple pairs with the same requirement type.
Otherwise please provide helpful response to the user
                """,
            },
            {
                "role": "user",
                "content": input_str,
            },
        ],
        model="gpt-4o-mini",
    )
    output = chat_completion.choices[0].message.content
    print(f"model output {output}")

    message_out = ""

    def set_requirement(output):
        nonlocal message_out
        if ":" not in output:
            return
        parts = output.split(":")
        req = parts[0].strip()
        content = output[len(req) + 1 :].strip()
        if req in REQUIREMENTS_KEYS:
            message_out = (
                message_out
                + f"\nThanks! Collected requirement: {req}\n{get_requirement_prompt()}"
            )
            requirements[req] = content

    if "|" in output:
        all_requirements = output.split("|")
        for req in all_requirements:
            set_requirement(req)
    else:
        set_requirement(output)

    return message_out.strip()


chat_history = []


def find_property(input_str):
    global requirements
    global chat_history

    chat_completion = client.chat.completions.create(
        messages=[
            {
                "role": "system",
                "content": f"You are a real estate agent looking for properties for your clients. Here are some properties that might be of interest to you: {json.dumps(props_core)}. Client is looking for properties that meet the following requirements: {requirements}",
            },
        ]
        + chat_history,
        model="gpt-4o-mini",
    )

    output = chat_completion.choices[0].message.content

    chat_history.append(
        {
            "role": "assistant",
            "content": output,
        }
    )

    return output


def process_chat_message(message, history):
    global chat_history
    output = ""
    if len(requirements) < len(REQUIREMENTS_KEYS):
        output = get_requirements(message)
        if len(requirements) == len(REQUIREMENTS_KEYS):
            output += "\n" + find_property(message)
    else:
        chat_history.append(
            {
                "role": "user",
                "content": message,
            }
        )
        output = find_property(message)
    return output


demo = gr.ChatInterface(
    fn=process_chat_message,
    chatbot=gr.Chatbot(value=[[None, get_requirement_prompt()]]),
    examples=[
        "I want a house near Houston",
        "I have two kids, what type of house would I need?",
        "I have a budget of 1 million dollars",
    ],
    title="Buyer Agent Bot",
)
demo.launch(share=True)