import json
import numpy as np
from PIL import Image
import subprocess
import sys
import logging
import gradio as gr
from typing import Tuple, Dict, Optional
import tempfile
import os

def install(package):
    subprocess.check_call([sys.executable, "-m", "pip", "install", package])

# 依存関係のインストール
install("opencv-python-headless")
install("moondream")
install("openai")
import cv2
import moondream as md
from openai import OpenAI
class PrivacyProtector:
    def __init__(self, moondream_api_key, deepseek_api_key):
        self.moon_model = md.vl(api_key=moondream_api_key)
        self.deepseek_client = OpenAI(
            api_key=deepseek_api_key,
            base_url="https://api.deepseek.com"
        )

    def analyze_risk(self, image_path):
        """画像のリスク分析を行う"""
        try:
            pil_image = Image.open(image_path)
            cv_image = cv2.imread(image_path)
            if cv_image is None:
                raise ValueError("画像の読み込みに失敗しました")

            encoded_image = self.moon_model.encode_image(pil_image)
            caption = self.moon_model.caption(encoded_image)["caption"]
            
            analysis_prompt = f"""
            以下の画像説明を基に個人情報漏洩リスクを分析し、厳密にJSON形式で返答してください：
            {{
                "risk_level": "high|medium|low",
                "risk_reason": "リスクの具体的理由",
                "objects_to_remove": ["消去すべきオブジェクトリスト"]
            }}
            
            画像説明: {caption}
            """
            
            response = self.deepseek_client.chat.completions.create(
                model="deepseek-chat",
                messages=[
                    {"role": "system", "content": "あなたは優秀なセキュリティ分析AIです。"},
                    {"role": "user", "content": analysis_prompt}
                ],
                temperature=0.3,
                response_format={"type": "json_object"}
            )
            
            result = json.loads(response.choices[0].message.content)
            return pil_image, cv_image, result
        
        except Exception as e:
            logging.error(f"リスク分析エラー: {str(e)}")
            raise

    def remove_objects(self, pil_image, cv_image, objects_to_remove):
        """オブジェクト検出と消去処理"""
        try:
            mask = np.zeros(cv_image.shape[:2], dtype=np.uint8)
            h, w = cv_image.shape[:2]
            
            for obj_name in objects_to_remove:
                detection = self.moon_model.detect(pil_image, obj_name)
                for obj in detection["objects"]:
                    x_min = int(obj['x_min'] * w)
                    y_min = int(obj['y_min'] * h)
                    x_max = int(obj['x_max'] * w)
                    y_max = int(obj['y_max'] * h)
                    cv2.rectangle(mask, (x_min, y_min), (x_max, y_max), 255, -1)
            
            inpainted = cv2.inpaint(cv_image, mask, 3, cv2.INPAINT_TELEA)
            return inpainted
        
        except Exception as e:
            logging.error(f"オブジェクト削除エラー: {str(e)}")
            raise

    def process_image(self, image_path):
        """画像処理フロー全体"""
        try:
            pil_img, cv_img, result = self.analyze_risk(image_path)
            if result['risk_level'] != 'low':
                cleaned = self.remove_objects(pil_img, cv_img, result['objects_to_remove'])
                return cleaned, result
            return cv_img, result
        
        except Exception as e:
            logging.error(f"画像処理エラー: {str(e)}")
            raise

def gradio_process(input_image):
    try:
        # 一時ファイルの処理
        with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as tmp_file:
            input_path = tmp_file.name
            cv2.imwrite(input_path, cv2.cvtColor(input_image, cv2.COLOR_RGB2BGR))
        
        # 環境変数の安全な取得
        protector = PrivacyProtector(
            moondream_api_key=os.environ.get('MOONDREAM'),
            deepseek_api_key=os.environ.get('DEEPSEEK')
        )
        
        # 画像処理実行
        processed_image, result = protector.process_image(input_path)
        
        # 結果の整形
        output_image = cv2.cvtColor(processed_image, cv2.COLOR_BGR2RGB)
        
        info_html = f"""
        <div style="padding:20px; background:#f0f0f0; border-radius:10px;">
            <h3>分析結果</h3>
            <p><strong>生成キャプション:</strong> {result.get('caption', '')}</p>
            <p><strong>リスクレベル:</strong> <span style="color:{'red' if result['risk_level'] == 'high' else 'orange' if result['risk_level'] == 'medium' else 'green'}">{result['risk_level'].upper()}</span></p>
            <p><strong>理由:</strong> {result['risk_reason']}</p>
            <p><strong>消去対象:</strong> {', '.join(result['objects_to_remove'])}</p>
        </div>
        """
        
        return input_image, output_image, info_html
    
    except Exception as e:
        error_msg = f"処理中にエラーが発生しました: {str(e)}"
        return input_image, None, f'<div style="color:red;">{error_msg}</div>'
    
    finally:
        # 一時ファイルのクリーンアップ
        if 'input_path' in locals() and os.path.exists(input_path):
            os.remove(input_path)
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🛡️ プライバシー保護画像処理ツール")
    
    with gr.Row():
        with gr.Column(scale=2):
            input_img = gr.Image(label="入力画像", type="numpy")
            submit_btn = gr.Button("画像を分析・処理", variant="primary")
        
        with gr.Column(scale=2):
            output_img = gr.Image(label="処理後の画像")
            info_output = gr.HTML(label="分析結果")
    
    with gr.Accordion("処理の詳細", open=False):
        gr.Markdown("""
        **処理フロー:**  
        1. 画像キャプション生成 (Moondream)  
        2. リスク分析 (DeepSeek)  
        3. オブジェクト検出・消去 (Moondream + OpenCV)  
        """)
    
    submit_btn.click(
        fn=gradio_process,
        inputs=input_img,
        outputs=[input_img, output_img, info_output]
    )

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