import sys, os
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

import gradio as gr
from rdkit import Chem
from rdkit.Chem import Draw
from ONNX0630 import main as predict_smiles

import tempfile
import logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

def visualize_molecule(smiles):

    try:
        mol = Chem.MolFromSmiles(smiles)
        if mol is None:
            return None
        img = Draw.MolToImage(mol)
        return img
    except:
        return None

def predict(image):
    try:
        if image is None:
            logger.error("上传的图片为空")
            return "错误: 请上传有效的图片文件"

        # 验证工作目录
        work_dir = "/home/user/app"
        if not os.path.isdir(work_dir):
            logger.error(f"工作目录不存在或不是目录: {work_dir}")
            return f"错误: 工作目录无效 - {work_dir}"

        with tempfile.NamedTemporaryFile(suffix=".png", dir=work_dir, delete=False) as temp_file:
            temp_path = os.path.abspath(temp_file.name)
            logger.info(f"保存图片到: {temp_path}, 类型: {type(temp_path)}")
            image.save(temp_path)

        if not os.path.exists(temp_path):
            logger.error(f"图片文件未正确保存: {temp_path}")
            return f"错误: 图片文件未正确保存 - {temp_path}"

        logger.info(f"传递给 main 的路径: {temp_path}")
        smiles = predict_smiles([temp_path])
        mol_image = visualize_molecule(smiles)

        try:
            os.remove(temp_path)
            logger.info(f"已删除临时文件: {temp_path}")
        except Exception as e:
            logger.warning(f"删除临时文件失败: {e}")

        if mol_image is None:
            return None, smiles
            
        return mol_image, smiles
        
    except FileNotFoundError as e:
        logger.error(f"文件路径错误: {e}")
        return f"错误: 文件路径问题 - {str(e)}"
    except IsADirectoryError as e:
        logger.error(f"路径是目录: {e}")
        return f"错误: 路径是目录 - {str(e)}"
    except PermissionError as e:
        logger.error(f"权限错误: {e}")
        return f"错误: 文件权限问题 - {str(e)}"
    except Exception as e:
        logger.error(f"处理错误: {e}")
        return f"错误: {str(e)}"

# 定义 Gradio 界面
iface = gr.Interface(
    fn=predict,
    inputs=gr.Image(label="Upload molecular image", type="pil"),
    outputs=[
        gr.Image(label="Mol Image"),
        gr.Text(label="SMILES"),
    ],
    title="Optical Chemical Structure Recognition",
    description="Convert a molecular image into SMILES.<br> ",
    examples=[
        ["example.png"]
    ]
)

if __name__ == "__main__":
    iface.launch(server_name="0.0.0.0", server_port=7860, share=False)