import gradio as gr
from openai import OpenAI
import requests
from PIL import Image
from io import BytesIO
import base64
import chess
import io
import chess.pgn
import chess.svg
from cairosvg import svg2png



def pgn_string_to_base64_image(pgn_string):
    pgn_io = io.StringIO(pgn_string)
    game = chess.pgn.read_game(pgn_io)
    board = game.end().board()
    svg_string = chess.svg.board(board=board)
    png_bytes = svg2png(bytestring=svg_string.encode('utf-8'))
    base64_image = base64.b64encode(png_bytes).decode('utf-8')
    image_url = f"data:image/jpeg;base64,{base64_image}"
    image_stream = BytesIO(png_bytes)
    image = Image.open(image_stream)
    return image_url, image


def process_text(api_key, example, question):
    if example:
        question=example
    client = OpenAI(api_key=api_key)

    if 'checkmate' in question:
        pgn_string = question.split('\n\n')[1]        
        image_url, image = pgn_string_to_base64_image(pgn_string)
    else:
        image_response = client.images.generate(model="dall-e-3", prompt=question,size="1024x1024",quality="standard",n=1,)
        image_url = image_response.data[0].url
        response = requests.get(image_url)
        if response.status_code == 200:
            image = Image.open(BytesIO(response.content))
            
    chat_response = client.chat.completions.create(model="gpt-4-vision-preview", messages=[{"role": "user","content": [{"type": "text", "text": question}, {"type": "image_url","image_url": {"url": image_url,},},],}],max_tokens=300,)
    answer = chat_response.choices[0].message.content
    
    return image, answer

title_html = """
<a href='https://github.com/GraphPKU/CoI' target='coi'>Chain of Images for Intuitively Reasoning</a>
"""

demo = gr.Interface(
    fn=process_text,
    inputs=[
        gr.Textbox(label="Your API Key", type="password"), 
        gr.Radio(["A group of people are crowded around in a living room talking to one another. A man in the foreground introduces two individuals one appears to be a regular human male the other appears to be an animal. What is unusual about this description?", \
            "A woman is waiting to get on the elevator. But the people in the elevator are on fire. Where can this event take place?",
            "In the following chess position, find a checkmate-in-one move.\n\n1. d4 d5 2. e3 Nf6 3. c4 Bf5 4. Bd3 e6 5. Bxf5 exf5 6. cxd5 Qxd5 7. Nf3 Nbd7 8. a3 O-O-O 9. Nc3 Qa5 10. Bd2 Be7 11. b4 Qa6 12. Qe2 Qxe2+ 13. Nxe2 Ne4 14. O-O g5 15. Rfd1 f6 16. d5 Bd6 17. Nfd4 Nxd2 18. Rxd2 Be5 19. Rad1 f4 20. exf4 gxf4 21. Nf3 Rdg8 22. Kf1 h5 23. h3 Bd6 24. Rd4 Ne5 25. Nxe5 fxe5 26. Re4 Rg6 27. Rd3 Rhg8 28. g3 fxg3 29. Rxg3 Rxg3 30. Nxg3 Rf8 31. Nxh5 Rf3 32. Ng3 Rxa3 33. Nf5 Ra1+ 34. Kg2 Ra6 35. Nxd6+ Rxd6 36. Rxe5 Kd7 37. Rh5 Rg6+ 38. Kf1 Kd6 39. Rf5 c6 40. dxc6 Kxc6 41. Rc5+ Kd6 42. Rc1 b6 43. Rd1+ Kc6 44. Rc1+ Kb5 45. Rb1 a5 46. bxa5+ Kxa5 47. Ra1+ Kb4 48. Rb1+ Kc3 49. Re1 b5 50. h4 b4 51. h5 Rg5 52. h6 Rh5 53. Re6 b3 54. Rc6+ Kb2 55. Rb6 Kb1 56. Ke2 b2 57. Kd2 Ka2 58. Kc3 b1=Q 59. Rxb1 Kxb1 60. Kd3 Rxh6 61. Ke3 Re6+ 62. Kf4 Kc2 63. Kf5 Re8 64. f4 Kd3 65. Kf6 Ke4 66. f5 Rb8 67. Ke6 Rb6+ 68. Kf7 Kxf5 69. Ke7 Re6+ 70. Kd7 Ke5 71. Kc7 Kd5 72. Kd7 Ke5 73. Kc7 Kd4 74. Kd7 Kd5 75. Kc7 Rd6 76. Kb7 Rc6 77. Kb8 Kc4 78. Kb7 Kb5 79. Ka7 Rb6 80. Ka8 Kc6 81. Ka7 Kc7 82. Ka8",
            ], label="Example Question"),
        gr.Textbox(label="Question")
    ],
    outputs=[
        gr.Image(type="pil", label="Generated Image", image_mode="fixed", width=768, height=768),
        gr.Textbox(label="Answer")],
    title=title_html
)

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