import gradio as gr
from ultralytics import YOLO
from PIL import Image
from ultralytics.utils.plotting import Annotator, colors
import glob
from database_center import db_transaction
from cloudhands import CloudHandsPayment
import uuid
import os 

payment_key=os.getenv('Payment_Key')
# Load model and data
model = YOLO('Dental_model.pt')
pic_files = glob.glob('*.jpg')
names = model.model.names

# ...existing code...
def detect_objects(state, image):
    chpay = state['chpay']
    payment_result = None

    # attempt to charge only if we have a token
    if state.get('token'):
        try:
            payment_result = chpay.charge(charge=0.2)
        except Exception as e:
            print(f"Error charging payment: {e}")
            payment_result = None
    else:
        print("No payment token. Please authorize first.")

    # if no successful payment, return empty outputs (Gradio expects two outputs)
    if not getattr(payment_result, "transaction_id", None):
        gr.Info('Transaction Failed')
        return None, None,

    if image is None:
        return None, None
    gr.Info('Transaction completed')
    db_transaction.add({
            'id':str(uuid.uuid4()),
            'app':'Dental Analysis',
            'transaction-id':payment_result.transaction_id,
            'price':0.1

            })
    image1 = image.copy()
    detection_results = model.predict(image)
    classes = detection_results[0].boxes.cls.cpu().tolist()
    boxes = detection_results[0].boxes.xyxy.cpu()
    annotator = Annotator(image, line_width=3)
    annotator1 = Annotator(image1, line_width=3)

    for box, cls in zip(boxes, classes):
        annotator.box_label(box, label=names[int(cls)], color=colors(int(cls)))
        annotator1.box_label(box, label=None, color=colors(int(cls)))

    return Image.fromarray(annotator.result()), Image.fromarray(annotator1.result()),'Payment succeed'
# ...existing code...

# ...existing code...
def Authorise_link(state):
    chpay = state['chpay']
    url = chpay.get_authorization_url()
    # return an HTML anchor and attempt to open a new tab (browser may block auto-open)
    html = f'<a href="{url}" target="_blank" rel="noopener noreferrer">Open authorization in a new tab</a>'
    html += f'<script>window.open("{url}", "_blank");</script>'
    return html
# ...existing code...

def set_api_key(state, payment_api_key):
    chpay = state['chpay']
    token = chpay.exchange_code_for_token(payment_api_key)
    state['token'] = token
    return state

# Gradio Blocks App
with gr.Blocks() as demo:
    gr.Markdown("## Dental Analysis")
    gr.Markdown("Analyze your Dental XRAY image with our AI object Detection model")
    payment_state = gr.State({'chpay':CloudHandsPayment(author_key=payment_key), 'token':None, 'transaction_id':None})

    with gr.Row():
        with gr.Column():
            apiKey = gr.Button("Get_api_key")
            # add an HTML output to receive the returned link/script
            auth_link = gr.HTML("")
            apiKey.click(fn=Authorise_link, inputs=payment_state, outputs=auth_link)
            payment_api_key = gr.Textbox(label="Authorization token", interactive=True, type="password")
            authorise_button=gr.Button("Authorise Payment")
            authorise_button.click(fn=set_api_key, inputs=(payment_state, payment_api_key), outputs=payment_state)
            image_input = gr.Image(type="pil", label="Upload Image")
            run_button = gr.Button("Run Detection")
            example_images = gr.Examples(
                examples=pic_files,
                inputs=image_input,
                label="Examples"
            )
        with gr.Column():
            image_output_1 = gr.Image(type="pil", label="Dental Analysis")
            image_output_2 = gr.Image(type="pil", label="Without Labels")
        run_button.click(detect_objects,inputs=[payment_state,image_input],outputs=[image_output_1,image_output_2])

if __name__ == "__main__":
    demo.launch()