# app.py
import gradio as gr
import requests
import base64
import os
import openai
from datetime import datetime
import random
import re
import markdown

# 设置 API 密钥（通过 Hugging Face Space Secrets 配置）
API_KEY = os.getenv("BAIDU_API_KEY")
SECRET_KEY = os.getenv("BAIDU_SECRET_KEY")
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")

# 百度 OCR 接口
def get_access_token(api_key, secret_key):
    url = "https://aip.baidubce.com/oauth/2.0/token"
    params = {
        "grant_type": "client_credentials",
        "client_id": api_key,
        "client_secret": secret_key
    }
    response = requests.post(url, params=params)
    return response.json().get("access_token")

def ocr_image(image_path, access_token):
    with open(image_path, "rb") as f:
        img_data = base64.b64encode(f.read()).decode()
    
    url = f"https://aip.baidubce.com/rest/2.0/ocr/v1/handwriting?access_token={access_token}"
    headers = {"Content-Type": "application/x-www-form-urlencoded"}
    data = {
        "image": img_data,
        "language_type": "ENG"
    }
    response = requests.post(url, headers=headers, data=data)
    return response.json()

# GPT 处理函数
def process_with_gpt(prompt):
    client = openai.OpenAI(
        api_key=OPENAI_API_KEY,
        base_url="https://api.openai.com/v1"
    )
    
    completion = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[{"role": "user", "content": prompt}],
        temperature=0.7
    )
    return completion.choices[0].message.content

# 主处理流程
def process_essay(image):
    # 生成文件名
    now = datetime.now().strftime("%Y%m%d%H%M%S")
    code = f"{random.randint(0000,9999):04}"
    file_name = "#" + str(code) + "-" + now
    
    # 临时保存上传的文件
    temp_path = f"temp_{code}.jpg"
    with open(temp_path, "wb") as f:
        f.write(image.read())
    
    try:
        # OCR 识别
        access_token = get_access_token(API_KEY, SECRET_KEY)
        result = ocr_image(temp_path, access_token)
        
        if "words_result" not in result:
            return "OCR 识别失败，请检查图片内容"
            
        essay_text = "\n".join([item["words"] for item in result["words_result"]])
        
        # 格式整理
        format_prompt = (
            "请帮我整理下面的英语作文文本格式，只整理英文正文部分（忽略英文正文外所有部分），"
            "保证原汁原味（明显错误空格换行、乱码、非常用字符比如☰需要改正除外），"
            "出现的拼写错误也不要帮助改正：\n\n" + essay_text
        )
        revised_text = process_with_gpt(format_prompt)
        
        # 语法纠正
        correction_prompt = (
            "请帮我把下面的英语作文的语法错误改正，输出改正后的文章（只改错误和不流畅之处）,"
            "请参照下面的格式要求\n"
            "格式要求：原文修改部分用()括起来，修改的部分用[]括起来，修改单词括单词、"
            "修改短语括短语、修改句子括句子，括的部分精准一些，能反映问题\n"
            "示例：The (rabbish) [rubbish] thrown by visitors has piled up and its lush (verdure no longer flourish) [verdure no longer flourishes] as it (did once) [once did].\n\n"
            "需要批改的英语习作：\n" + revised_text
        )
        corrected_text = process_with_gpt(correction_prompt)
        
        # 批改意见
        assessment_prompt = """
        下面是一份已经批改过的英语作文（改正了一些语法错误和不流畅不准确之处），其中小括号表示原文的错误，
        中括号表示原文的修改部分。请你根据修改的批注，按照如下格式给出批改意见：
        > 英文有错误的原文1
        - 错误1和解决方案
        - 错误2和解决方案，后面的以此类推
        示例：
        > The rabbish thrown by visitors has piled up and its lush verdure no longer flourish as it did once.
        - rabbish：拼写错误
        - verdure no longer flourish：动词单复数错误，其中verdure是单数，动词使用第三人称单数
        - as it did once：语序错误
        
        实际批改内容：
        """ + corrected_text
        
        assessment_text = process_with_gpt(assessment_prompt)
        
        # 评分系统
        rate_prompt = """
        你是资深英语写作批改专家，按以下三个维度评分（F/B-/B/B+/A-/A/A+等级）：
        ## 语言通顺度(Fluency)&可读性与风格(Readability & Style)
        ## 上下文连贯度(Coherence)
        ## 词汇多样性(Lexical Resource)&语法准确性(Grammatical Accuracy)
        要求：1. 给出评分等级 2. 详细分析原因 3. 提出改进方向
        示例格式：
        ## 语言通顺度 B+
        - 优点说明...
        - 缺点说明...
        - 改进建议...
        
        需要评分的作文：
        """ + revised_text
        
        rate_text = process_with_gpt(rate_prompt)
        
        # 优秀范文
        rewrite_prompt = """
        请使用优秀的英语表达重写下面这篇英文习作，要展现优秀的词汇和语法，使用地道的表达方式，
        使用多样化的句式、短语和词汇，加粗可供学习的部分。
        原文：
        """ + revised_text
        
        perfect_text = process_with_gpt(rewrite_prompt)
        
        # 推荐表达
        theme_prompt = """
        请分析作文内容并给出主题相关的优秀表达：
        1. 推测作文主题
        2. 分析具体内容
        3. 提供相关优秀表达（包含句子、句型、短语、词汇）
        
        需要分析的作文：
        """ + revised_text
        
        theme_text = process_with_gpt(theme_prompt)
        theme_text = theme_text.split("## 可用表达")[1] if "## 可用表达" in theme_text else ""
        
        # HTML 生成
        def highlight_brackets(text):
            text = re.sub(r'\[([^\[\]]+)\]', r'<span style="background:#d4f7d4">\1</span>', text)
            text = re.sub(r'\(([^\(\)]+)\)', r'<span style="background:#ffd6d6">\1</span>', text)
            return text.replace('\n', '<br>')
        
        # 图片 Base64 编码
        with open(temp_path, "rb") as img_f:
            img_b64 = base64.b64encode(img_f.read()).decode()
            img_html = f'<img src="data:image/jpeg;base64,{img_b64}" style="width:100%;max-height:600px;object-fit:contain;margin:20px 0">'
        
        # 组合最终 HTML
        final_html = f"""
        <html>
        <style>
            body {{ font-family: Arial, sans-serif; line-height: 1.6; padding: 20px; }}
            .section {{ margin-bottom: 30px; }}
            code {{ background: #f5f5f5; padding: 2px 4px; border-radius: 4px; }}
            pre {{ background: #f5f5f5; padding: 10px; overflow-x: auto; }}
            blockquote {{ border-left: 4px solid #ddd; padding-left: 12px; color: #666; }}
        </style>
        <body>
            <h1>作文批改结果 #{code}</h1>
            {img_html}
            
            <div class="section">
                <h2>批改结果</h2>
                <div style="font-family:Consolas,monospace;">{highlight_brackets(corrected_text)}</div>
            </div>
            
            <div class="section">
                <h2>批改意见</h2>
                <div style="font-family:Consolas,monospace;">{markdown.markdown(assessment_text)}</div>
            </div>
            
            <div class="section">
                <h2>评分分析</h2>
                <div style="font-family:Consolas,monospace;">{markdown.markdown(rate_text)}</div>
            </div>
            
            <div class="section">
                <h2>优秀范文</h2>
                <div style="font-family:Consolas,monospace;">{markdown.markdown(perfect_text)}</div>
            </div>
            
            <div class="section">
                <h2>推荐表达</h2>
                <div style="font-family:Consolas,monospace;">{markdown.markdown(theme_text)}</div>
            </div>
        </body>
        </html>
        """
        
        return final_html
        
    finally:
        # 清理临时文件
        if os.path.exists(temp_path):
            os.remove(temp_path)

# 创建 Gradio 界面
def create_interface():
    description = """
    英语作文批改助手
    上传一张包含英语作文的手写图片，系统将自动进行文字识别并提供详细的批改建议、评分和优秀范文。
    """
    
    interface = gr.Interface(
        fn=process_essay,
        inputs=gr.File(label="上传作文图片（JPG/PNG）", type="file"),
        outputs=gr.HTML(label="批改结果"),
        title="英语作文智能批改",
        description=description,
        #examples=[["example.jpg"]],
        #cache_examples=False,
        theme="default"
    )
    
    return interface

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