import json
import gradio as gr
from dotenv import load_dotenv
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
from openai import BadRequestError

load_dotenv('.env', override=True)

def extract_free_fn(gallery, provider, schema):
    img_paths = [item[0] for item in gallery]
    # 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
            with open('read_schema_free.py', 'w') as f:
                f.write(full_schema)
            exec(import_part + 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:
        free_attributes = extract_info(img_paths, provider, Garment)
        free_attributes = free_attributes.model_dump_json()
    except RefusalError as e:
        return e.message, {}, []
    except BadRequestError as e:
        return e.message, {}, []

    return f'OK - provider: {provider}', free_attributes, [1] * 100

def follow_schema_fn(free_attributes, provider, probs, schema):
    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:
        schema_attributes = follow_structure(free_attributes, provider, Garment)
        schema_attributes = schema_attributes.model_dump_json()
    except RefusalError as e:
        return e.message, {}
    except BadRequestError as e:
        return e.message, {}
    
    return f'OK - 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):
            # gr.Markdown('''<div style="text-align: left; font-size: 18px;">Upload images of a garment and click Extract Freestyle</div>''')
            gallery = gr.Gallery(label='Images', type='filepath')
            probs_text = gr.Textbox(label='Probabilities', placeholder='Probabilities', visible=False)
            provider = gr.Dropdown(label='Provider', choices=['openai', 'anthropic'], value='openai')
            button_extract = gr.Button(value='Extract Freestyle')
        with gr.Column(scale=4):
            # gr.Markdown('''<div style="text-align: left; font-size: 18px;">Click Follow Schema to convert the freestyle output to predefined schema</div>''')
            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():
        # free_schema = gr.Dataframe(row_count=(4, 'dynamic'), label='Freestyle Schema', col_count=(4, 'dynamic'), interactive=True)
        # schema = gr.Dataframe(row_count=(4, 'dynamic'), label='Schema', col_count=(4, 'dynamic'), interactive=True)
        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, provider, free_schema], [status_free_text, free_attributes, probs_text])
    button_follow.click(follow_schema_fn, [free_attributes, provider, probs_text, structured_schema], [status_structured_text, schema_attributes])

# demo.launch(server_name="0.0.0.0", server_port=7688, share=False, ssl_verify=True, ssl_certfile=None, ssl_keyfile=None, auth=None)
# demo.launch(server_name="0.0.0.0", server_port=7860, share=False, ssl_verify=False, ssl_certfile="server.crt", ssl_keyfile="server.key", auth=None)
demo.launch(server_name="0.0.0.0", server_port=7861)