import json

import gradio as gr
from dotenv import load_dotenv
from openai import BadRequestError

from llm_api.exceptions import RefusalError

# from llm_api.langchain_api import extract_info, follow_structure
from llm_api.native_api import extract_info, follow_structure

load_dotenv(".env", override=True)
garment_categories = [
    "Upper Garment",
    "Lower Garment",
    "Hat",
    "Luggage",
    "Handbag",
    "Makeup",
    "Chair",
    "Carpet",
    "Watch",
    "Cookware",
    "Tableware",
    "Sleepware",
    "BathroomItem",
]
schema_categories = [
    "UpperGarmentProperties",
    "LowerGarmentProperties",
    "HatProperties",
    "LuggageProperties",
    "HandbagProperties",
    "MakeupProperties",
    "ChairProperties",
    "CarpetProperties",
    "WatchProperties",
    "CookwareProperties",
    "TablewareProperties",
    "SleepwareProperties",
    "BathroomItemProperties",
]


def extract_free_fn(gallery, category, provider, schema, known_data):
    img_paths = [item[0] for item in gallery]
    category_schema = schema_categories[garment_categories.index(category)]
    # write schema to file before calling extract_free_properties
    try:
        if (
            schema is not None
            and isinstance(schema, str)
            and "import"
            not in schema.replace("'import'", "").replace("'important'", "")
        ):
            with open("schema_free.py", "r") as f:
                default_schema = f.read()
            import_part = default_schema.split("#####")[0]
            print("Executing freestyle schema...")
            full_schema = import_part + schema

            # dynamically change the last line of the schema to include the category schema
            if category_schema in schema_categories:
                print("Dynamic schema activated")
                schema_lines = full_schema.split("\n")
                schema_lines[-1] = f"    properties: {category_schema}"
                full_schema = "\n".join(schema_lines)

            with open("read_schema_free.py", "w") as f:
                f.write(full_schema)
            exec(full_schema, globals())
            print("Executed freestyle schema")
        else:
            print("SCHEMA ELSE")
            raise Exception("Invalid schema")
    except Exception as e:
        print(e)
        return (
            "Invalid schema! Please check the schema or Reload this page to restore default settings",
            {},
            [],
        )

    try:
        status_code, free_attributes = extract_info(img_paths, provider, Garment, known_data=known_data)
        free_attributes = free_attributes.model_dump_json()
    except RefusalError as e:
        return e.message, {}, []
    except BadRequestError as e:
        return e.message, {}, []

    return f"status: {status_code} - provider: {provider}", free_attributes, [1] * 100


def follow_schema_fn(free_attributes, provider, probs, schema, known_data):
    try:
        if (
            schema is not None
            and isinstance(schema, str)
            and "import"
            not in schema.replace("'import'", "").replace("'important'", "")
        ):
            with open("schema_structured.py", "r") as f:
                default_schema = f.read()
            import_part = default_schema.split("#####")[0]
            print("executing structured schema\n")
            full_schema = import_part + schema
            with open("read_schema_structured.py", "w") as f:
                f.write(full_schema)
            exec(full_schema, globals())
            print("executed structured schema")
        else:
            print("SCHEMA ELSE")
            raise Exception("Invalid schema")
    except:
        return (
            "Invalid schema! Please check the schema or Reload this page to restore default settings",
            {},
        )

    try:
        status_code, schema_attributes = follow_structure(
            free_attributes, provider, Garment, known_data=known_data
        )
        schema_attributes = schema_attributes.model_dump_json()
    except RefusalError as e:
        return e.message, {}
    except BadRequestError as e:
        return e.message, {}

    return f"status {status_code} - provider: {provider}", schema_attributes


with gr.Blocks(title="Internal Demo for Attribution") as demo:
    with gr.Row():
        with gr.Column(scale=12):
            gr.Markdown(
                """<div style="text-align: center; font-size: 24px;"><strong>Internal Demo for Attribution</strong></div>"""
            )
    with gr.Row():
        with gr.Column(scale=4):
            gallery = gr.Gallery(label='Upload images of your product here',type="filepath")
            probs_text = gr.Textbox(
                label="Probabilities", placeholder="Probabilities", visible=False
            )
            category = gr.Dropdown(
                label="Category",
                choices=garment_categories,
                value=garment_categories[0],
            )
            provider = gr.Dropdown(
                label="Provider", choices=["openai", "anthropic"], value="anthropic"
            )
            known_data = gr.TextArea(
                label="Insert known attributes of the product here (Optional)", placeholder="", visible=True
            )
            button_extract = gr.Button(value="Extract Freestyle")
        with gr.Column(scale=4):
            status_free_text = gr.Textbox(label="Request Status", placeholder="")
            free_attributes = gr.JSON(label="Freestyle Attribute")
            button_follow = gr.Button(value="Follow Schema")
        with gr.Column(scale=4):
            status_structured_text = gr.Textbox(label="Request Status", placeholder="")
            schema_attributes = gr.JSON(label="Schema Attributes")
    with gr.Row():
        with open("schema_free.py", "r") as f:
            schema_free_text = f.read()
        schema_free_text = schema_free_text.split("#####")[1]
        with open("schema_structured.py", "r") as f:
            schema_structured_text = f.read()
        schema_structured_text = schema_structured_text.split("#####")[1]
        free_schema = gr.TextArea(
            label="Freestyle Schema",
            placeholder="Freestyle Schema",
            visible=True,
            value=schema_free_text,
            interactive=True,
        )
        structured_schema = gr.TextArea(
            label="Structured Schema",
            placeholder="Schema",
            visible=True,
            value=schema_structured_text,
            interactive=True,
        )

    button_extract.click(
        extract_free_fn,
        [gallery, category, provider, free_schema, known_data],
        [status_free_text, free_attributes, probs_text],
    )
    button_follow.click(
        follow_schema_fn,
        [free_attributes, provider, probs_text, structured_schema, known_data],
        [status_structured_text, schema_attributes],
    )

demo.launch()