import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
import PyPDF2
import requests
from bs4 import BeautifulSoup
import re
import warnings
import os
warnings.filterwarnings("ignore")

# ZeroGPU環境の検出
IS_ZEROGPU = os.environ.get("SPACE_ID") is not None and "zero-gpu" in os.environ.get("SPACE_ID", "").lower()

# spacesライブラリのインポート（ZeroGPU必須）
try:
    import spaces
    print("✅ spacesライブラリが正常にインポートされました")
    GPU_DECORATOR = spaces.GPU
except ImportError as e:
    print(f"⚠️  spacesライブラリのインポートに失敗: {e}")
    if IS_ZEROGPU:
        print("🚨 ZeroGPU環境ではspacesライブラリが必須です")
        raise ImportError("ZeroGPU環境でspacesライブラリが見つかりません。requirements.txtにspaces>=0.19.0を追加してください。")
    else:
        print("ℹ️  ローカル環境のため、spacesライブラリは不要です")
        # ローカル環境用のダミーデコレータ
        GPU_DECORATOR = lambda duration=None: lambda func: func

class TextSummarizer:
    def __init__(self):
        # 環境の検出
        self.is_zerogpu = IS_ZEROGPU
        print(f"実行環境: {'ZeroGPU' if self.is_zerogpu else 'Local/Standard'}")
        
        # デバイス設定
        if self.is_zerogpu:
            self.device = "cuda"  # ZeroGPUでは常にCUDA
        else:
            self.device = "cuda" if torch.cuda.is_available() else "cpu"
        
        print(f"使用デバイス: {self.device}")
        print(f"PyTorch バージョン: {torch.__version__}")
        
        # ZeroGPU以外でのバージョンチェック
        if not self.is_zerogpu:
            torch_version = torch.__version__.split('+')[0]
            try:
                major, minor = map(int, torch_version.split('.')[:2])
                if major < 2 or (major == 2 and minor < 6):
                    print("⚠️  警告: PyTorch v2.6未満です。セキュリティ脆弱性(CVE-2025-32434)のため、アップグレードを推奨します。")
            except ValueError:
                print("⚠️  PyTorchバージョンの解析に失敗しました")
        
        # モデル読み込み（ZeroGPU対応）
        self._load_model()
    
    def _load_model(self):
        """モデルの読み込み（環境に応じて最適化）"""
        try:
            print("モデルを読み込み中...")
            
            # ZeroGPU環境では軽量モデルを優先
            # if self.is_zerogpu:
            model_name = "tsmatz/mt5_summarize_japanese"
            # else:
            #     model_name = "facebook/bart-large-cnn"
            
            # モデル読み込み設定（use_safetensorsを削除）
            pipeline_kwargs = {
                "task": "summarization",
                "model": model_name,
                "device": 0 if self.device == "cuda" else -1,
                "framework": "pt",
                "trust_remote_code": False  # セキュリティ強化
            }
            
            # 基本的なパイプライン作成
            self.summarizer = pipeline(**pipeline_kwargs)
            print(f"✅ {model_name} の読み込み完了")
                
        except Exception as e:
            print(f"❌ メインモデル読み込みエラー: {e}")
            # 最軽量フォールバック
            try:
                print("最軽量フォールバックモデルを試行中...")
                self.summarizer = pipeline(
                    "summarization",
                    model="sshleifer/distilbart-cnn-6-6",
                    device=0 if self.device == "cuda" else -1,
                    trust_remote_code=False
                )
                print("✅ 最軽量モデルで読み込み完了")
            except Exception as e2:
                print(f"❌ 全てのモデル読み込みに失敗: {e2}")
                raise Exception(f"モデルの読み込みに失敗しました: {e2}")
    
    def clean_text(self, text):
        """テキストの前処理"""
        # 不要な文字や改行を整理
        text = re.sub(r'\n+', '\n', text)
        text = re.sub(r'\s+', ' ', text)
        text = text.strip()
        return text
    
    def chunk_text(self, text, max_length=1000):
        """長いテキストをチャンクに分割"""
        sentences = text.split('.')
        chunks = []
        current_chunk = ""
        
        for sentence in sentences:
            if len(current_chunk + sentence) < max_length:
                current_chunk += sentence + "."
            else:
                if current_chunk:
                    chunks.append(current_chunk.strip())
                current_chunk = sentence + "."
        
        if current_chunk:
            chunks.append(current_chunk.strip())
        
        return chunks
    
    def summarize_text(self, text, max_length=150, min_length=50):
        """テキストを要約（ZeroGPU対応）"""
        try:
            cleaned_text = self.clean_text(text)
            
            if len(cleaned_text) < 100:
                return "テキストが短すぎるため、要約できません。"
            
            # テキストが長い場合はチャンクに分割
            if len(cleaned_text) > 1000:
                chunks = self.chunk_text(cleaned_text)
                summaries = []
                
                for chunk in chunks:
                    try:
                        result = self.summarizer(
                            chunk,
                            max_length=max_length,
                            min_length=min_length,
                            do_sample=False
                        )
                        summaries.append(result[0]['summary_text'])
                    except Exception as e:
                        print(f"チャンク要約エラー: {e}")
                        continue
                
                # チャンクの要約を統合
                combined_summary = " ".join(summaries)
                if len(combined_summary) > max_length * 2:
                    # 再度要約
                    final_result = self.summarizer(
                        combined_summary,
                        max_length=max_length,
                        min_length=min_length,
                        do_sample=False
                    )
                    return final_result[0]['summary_text']
                else:
                    return combined_summary
            else:
                result = self.summarizer(
                    cleaned_text,
                    max_length=max_length,
                    min_length=min_length,
                    do_sample=False
                )
                return result[0]['summary_text']
                
        except Exception as e:
            return f"要約処理でエラーが発生しました: {str(e)}"
    
    def structure_summary(self, summary_text):
        """要約を構造化"""
        # 簡単な構造化ロジック（実際のプロジェクトではより高度な処理が必要）
        sentences = summary_text.split('.')
        
        structured_output = "## 📋 要約結果\n\n"
        
        if len(sentences) >= 3:
            structured_output += "### 🎯 主要ポイント\n"
            structured_output += f"- {sentences[0].strip()}\n\n"
            
            structured_output += "### 📊 詳細内容\n"
            for i, sentence in enumerate(sentences[1:-1], 1):
                if sentence.strip():
                    structured_output += f"{i}. {sentence.strip()}\n"
            
            if sentences[-1].strip():
                structured_output += f"\n### 💡 結論\n"
                structured_output += f"- {sentences[-1].strip()}\n"
        else:
            structured_output += f"### 📄 要約内容\n{summary_text}\n"
        
        return structured_output
    
    def extract_text_from_pdf(self, pdf_file):
        """PDFからテキストを抽出"""
        try:
            reader = PyPDF2.PdfReader(pdf_file)
            text = ""
            for page in reader.pages:
                text += page.extract_text() + "\n"
            return text
        except Exception as e:
            return f"PDFの読み込みでエラーが発生しました: {str(e)}"
    
    def extract_text_from_url(self, url):
        """Webサイトからテキストを抽出"""
        try:
            headers = {
                'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36'
            }
            response = requests.get(url, headers=headers, timeout=10)
            response.encoding = response.apparent_encoding
            
            soup = BeautifulSoup(response.text, 'html.parser')
            
            # 不要なタグを削除
            for tag in soup(['script', 'style', 'nav', 'header', 'footer']):
                tag.decompose()
            
            # テキストを抽出
            text = soup.get_text()
            return self.clean_text(text)
            
        except Exception as e:
            return f"Webサイトの読み込みでエラーが発生しました: {str(e)}"

# グローバルインスタンス（ZeroGPU環境では起動時に初期化が必要）
summarizer = None

def initialize_model():
    """モデルの初期化（ZeroGPU対応）"""
    global summarizer
    if summarizer is None:
        summarizer = TextSummarizer()
    return summarizer

@GPU_DECORATOR(duration=30)
def process_text_input(text, max_length, min_length):
    """テキスト入力の処理"""
    try:
        print(f"テキスト処理開始: {len(text) if text else 0}文字")
        
        if not text or not text.strip():
            return "## ⚠️ エラー\nテキストを入力してください。"
        
        # モデル初期化
        model = initialize_model()
        
        # 要約実行
        summary = model.summarize_text(text, int(max_length), int(min_length))
        result = model.structure_summary(summary)
        
        print("テキスト処理完了")
        return result
        
    except Exception as e:
        error_msg = f"## ❌ エラーが発生しました\n```\n{str(e)}\n```"
        print(f"テキスト処理エラー: {e}")
        return error_msg

@GPU_DECORATOR(duration=30)
def process_pdf_input(pdf_file, max_length, min_length):
    """PDF入力の処理"""
    try:
        print(f"PDF処理開始: {pdf_file}")
        
        if pdf_file is None:
            return "## ⚠️ エラー\nPDFファイルを選択してください。"
        
        # モデル初期化
        model = initialize_model()
        
        text = model.extract_text_from_pdf(pdf_file)
        if text.startswith("PDFの読み込みで"):
            return f"## ❌ エラー\n{text}"
        
        summary = model.summarize_text(text, int(max_length), int(min_length))
        result = model.structure_summary(summary)
        
        print("PDF処理完了")
        return result
        
    except Exception as e:
        error_msg = f"## ❌ エラーが発生しました\n```\n{str(e)}\n```"
        print(f"PDF処理エラー: {e}")
        return error_msg

@GPU_DECORATOR(duration=30)
def process_url_input(url, max_length, min_length):
    """URL入力の処理"""
    try:
        print(f"URL処理開始: {url}")
        
        if not url or not url.strip():
            return "## ⚠️ エラー\nURLを入力してください。"
        
        if not url.startswith(('http://', 'https://')):
            url = 'https://' + url
        
        # モデル初期化
        model = initialize_model()
        
        text = model.extract_text_from_url(url)
        if text.startswith("Webサイトの読み込みで"):
            return f"## ❌ エラー\n{text}"
        
        summary = model.summarize_text(text, int(max_length), int(min_length))
        result = model.structure_summary(summary)
        
        print("URL処理完了")
        return result
        
    except Exception as e:
        error_msg = f"## ❌ エラーが発生しました\n```\n{str(e)}\n```"
        print(f"URL処理エラー: {e}")
        return error_msg

# Gradioインターフェース作成
def create_interface():
    with gr.Blocks(title="🤖 ローカルLLM テキスト要約ツール", theme=gr.themes.Soft()) as app:
        gr.Markdown("""
        # 🤖 ローカルLLM テキスト要約ツール (ZeroGPU対応)
        
        このツールは、ローカルで動作するLLMを使用してテキストを要約し、構造化された形式で出力します。
        
        ## 🚀 環境対応
        - **ZeroGPU**: Hugging Face Spaces自動最適化
        - **ローカル環境**: 高性能モデル対応
        - **セキュリティ**: trust_remote_code=False設定
        
        ## 📝 対応入力形式
        - **テキスト直接入力**
        - **PDFファイル**
        - **Webサイト URL**
        """)
        
        # 要約設定
        with gr.Row():
            max_length = gr.Slider(
                minimum=50, maximum=500, value=150, step=10,
                label="最大要約長", info="要約の最大文字数"
            )
            min_length = gr.Slider(
                minimum=20, maximum=200, value=50, step=10,
                label="最小要約長", info="要約の最小文字数"
            )
        
        # タブインターフェース
        with gr.Tabs():
            # テキスト入力タブ
            with gr.TabItem("📝 テキスト入力"):
                with gr.Row():
                    with gr.Column():
                        text_input = gr.Textbox(
                            lines=10,
                            placeholder="要約したいテキストを入力してください...",
                            label="入力テキスト"
                        )
                        text_btn = gr.Button("🔍 要約実行", variant="primary")
                    
                    with gr.Column():
                        text_output = gr.Markdown(label="要約結果")
                
                text_btn.click(
                    fn=process_text_input,
                    inputs=[text_input, max_length, min_length],
                    outputs=text_output,
                    show_progress=True
                )
            
            # PDF入力タブ
            with gr.TabItem("📄 PDF入力"):
                with gr.Row():
                    with gr.Column():
                        pdf_input = gr.File(
                            file_types=[".pdf"],
                            label="PDFファイルを選択"
                        )
                        pdf_btn = gr.Button("🔍 PDF要約実行", variant="primary")
                    
                    with gr.Column():
                        pdf_output = gr.Markdown(label="要約結果")
                
                pdf_btn.click(
                    fn=process_pdf_input,
                    inputs=[pdf_input, max_length, min_length],
                    outputs=pdf_output,
                    show_progress=True
                )
            
            # URL入力タブ
            with gr.TabItem("🌐 Website URL"):
                with gr.Row():
                    with gr.Column():
                        url_input = gr.Textbox(
                            placeholder="https://example.com",
                            label="ウェブサイトURL"
                        )
                        url_btn = gr.Button("🔍 Web要約実行", variant="primary")
                    
                    with gr.Column():
                        url_output = gr.Markdown(label="要約結果")
                
                url_btn.click(
                    fn=process_url_input,
                    inputs=[url_input, max_length, min_length],
                    outputs=url_output,
                    show_progress=True
                )
        
        # 使用方法
        gr.Markdown("""
        ## 🔧 使用方法
        
        1. **要約設定**: 最大・最小要約長を調整
        2. **入力方法選択**: テキスト直接入力、PDFアップロード、URL入力から選択
        3. **実行**: 対応する実行ボタンをクリック
        4. **結果確認**: 構造化された要約結果を確認
        
        ## ⚙️ 技術仕様 (ZeroGPU対応)
        - **モデル**: DistilBART/BART (環境に応じて自動選択)
        - **ZeroGPU**: Hugging Face Spaces最適化
        - **セキュリティ**: trust_remote_code=False設定
        - **GPU加速**: 環境自動検出
        - **出力形式**: 構造化Markdown
        
        ## 🔧 環境別最適化
        - ZeroGPU: 軽量モデル自動選択
        - ローカル: 高性能モデル利用可能
        - 互換性: 幅広いTransformersバージョン対応
        """)
    
    return app

if __name__ == "__main__":
    # 環境情報表示
    print(f"""
    🚀 テキスト要約ツール起動 🚀
    
    実行環境: {'ZeroGPU (Hugging Face Spaces)' if IS_ZEROGPU else 'ローカル環境'}
    PyTorchバージョン: {torch.__version__}
    spacesライブラリ: {'✅ 利用可能' if 'spaces' in globals() else '❌ 未インストール'}
    
    """)
    
    if IS_ZEROGPU:
        print("✅ ZeroGPU環境で最適化済み")
        print("🔧 @spaces.GPU デコレータが適用されています")
    else:
        # ローカル環境でのセキュリティチェック
        try:
            torch_version = torch.__version__.split('+')[0]
            major, minor = map(int, torch_version.split('.')[:2])
            if major < 2 or (major == 2 and minor < 6):
                print("⚠️  セキュリティ警告: PyTorch v2.6未満")
                print("   推奨: pip install torch>=2.6.0")
            else:
                print("✅ PyTorchセキュリティ: OK")
        except ValueError:
            print("⚠️  PyTorchバージョン確認不能")
    
    # アプリケーション起動
    app = create_interface()
    
    if IS_ZEROGPU:
        # ZeroGPU環境用設定
        app.launch()
    else:
        # ローカル環境用設定
        app.launch(
            server_name="0.0.0.0",
            server_port=7860,
            share=True,
            debug=True
        )