"""
Script generation module using Google Gemini Pro

This module generates natural Japanese narration scripts based on:
- Extracted PDF data
- Video transcription
- Template constraints
"""

import os
import json
from typing import Dict, Any, List, Optional

try:
    import google.generativeai as genai
except ImportError:
    genai = None


class ScriptGenerator:
    """Generate narration scripts using Gemini Pro"""
    
    def __init__(self, api_key: Optional[str] = None):
        self.api_key = api_key or os.getenv("GEMINI_API_KEY")
        if not self.api_key:
            raise ValueError("GEMINI_API_KEY not found in environment")
        
        if genai is None:
            raise ImportError(
                "google-generativeai not installed. "
                "Run: pip install google-generativeai"
            )
        
        genai.configure(api_key=self.api_key)
        self.model = genai.GenerativeModel('gemini-1.5-pro')
    
    def generate(
        self, 
        pdf_data: Dict[str, Any], 
        video_data: Optional[Dict[str, Any]] = None,
        style: str = "friendly"
    ) -> Dict[str, Any]:
        """Generate narration script
        
        Args:
            pdf_data: Structured data from PDF extraction
            video_data: Optional video transcription data
            style: Narration style ('friendly', 'professional', 'casual')
            
        Returns:
            Dictionary containing:
            - segments: List of narration segments
            - total_duration: Estimated total duration in seconds
            - sync_points: Recommended chart/video sync points
        """
        print("   Generating script with Gemini Pro...")
        
        # Build the prompt
        prompt = self._build_prompt(pdf_data, style)
        
        try:
            # Call Gemini API
            response = self.model.generate_content(prompt)
            
            # Parse the response
            script_data = self._parse_response(response.text)
            
            print(f"   ✓ Generated {len(script_data['segments'])} script segments")
            return script_data
            
        except Exception as e:
            print(f"   ⚠️ Gemini API error: {e}")
            print("   Falling back to template-based script...")
            return self._generate_fallback_script(pdf_data)
    
    def _build_prompt(self, pdf_data: Dict, style: str) -> str:
        """Build the prompt for Gemini Pro"""
        client_name = pdf_data.get("client_name", "お客様")
        current_age = pdf_data.get("current_age", 35)
        retirement_age = pdf_data.get("retirement_age", 65)
        current_assets = pdf_data.get("current_assets", 5000)
        annual_income = pdf_data.get("annual_income", 600)
        annual_expenses = pdf_data.get("annual_expenses", 450)
        
        prompt = f"""
あなたはファイナンシャルプランナーのナレーション原稿作成アシスタントです。
以下の顧客情報に基づいて、3〜5分の解説動画用のナレーション原稿を作成してください。

【顧客情報】
- お名前: {client_name}様
- 現在年齢: {current_age}歳
- 退職予定年齢: {retirement_age}歳
- 現在の資産: {current_assets}万円
- 年収: {annual_income}万円
- 年間支出: {annual_expenses}万円

【原稿作成の指示】
1. 合計3〜5分（180〜300秒）の原稿を作成してください
2. 以下のセグメントに分割してください：
   - オープニング（挨拶と概要説明）
   - 現在の状況説明
   - 将来の資産推移の解説
   - 退職後の生活資金について
   - リスクと対策の提案
   - クロージング（まとめと安心感）

3. 各セグメントには以下を含めてください：
   - id: セグメント番号（1から開始）
   - text: ナレーション文（自然な日本語、敬語）
   - duration: 推定時間（秒）
   - type: シーンタイプ（"narration", "chart_intro", "chart_highlight", "transition"）
   - chartSync: チャート表示の指示（該当する場合のみ）

4. トーンは{style}（親しみやすく、でも専門的）にしてください

5. 出力はJSON形式で、以下の構造で返してください：
{{
  "segments": [
    {{
      "id": 1,
      "text": "ナレーション文",
      "duration": 3.5,
      "type": "narration",
      "chartSync": {{
        "chartId": "assetDepletion",
        "highlightYear": 2025
      }}
    }}
  ]
}}

必ずJSON形式のみを返してください。説明文は不要です。
"""
        return prompt
    
    def _parse_response(self, response_text: str) -> Dict[str, Any]:
        """Parse Gemini's JSON response"""
        # Remove markdown code blocks if present
        text = response_text.strip()
        if text.startswith("```json"):
            text = text[7:]
        if text.startswith("```"):
            text = text[3:]
        if text.endswith("```"):
            text = text[:-3]
        text = text.strip()
        
        # Parse JSON
        data = json.loads(text)
        
        # Calculate total duration
        total_duration = sum(seg.get("duration", 0) for seg in data["segments"])
        
        return {
            "segments": data["segments"],
            "total_duration": total_duration,
        }
    
    def _generate_fallback_script(self, pdf_data: Dict) -> Dict[str, Any]:
        """Generate a template-based script as fallback"""
        client_name = pdf_data.get("client_name", "お客様")
        current_age = pdf_data.get("current_age", 35)
        retirement_age = pdf_data.get("retirement_age", 65)
        
        segments = [
            {
                "id": 1,
                "text": f"{client_name}様、こんにちは。これから、あなたの人生設計について、詳しくご説明させていただきます。",
                "duration": 5.0,
                "type": "narration",
            },
            {
                "id": 2,
                "text": f"現在{current_age}歳の{client_name}様の資産状況を見てみましょう。",
                "duration": 3.5,
                "type": "chart_intro",
                "chartSync": {
                    "chartId": "cashFlow",
                },
            },
            {
                "id": 3,
                "text": f"{retirement_age}歳での退職を想定した場合、将来の資産推移はこのようになります。",
                "duration": 4.0,
                "type": "chart_highlight",
                "chartSync": {
                    "chartId": "assetDepletion",
                },
            },
            {
                "id": 4,
                "text": "退職後の生活資金については、計画的な準備が大切です。",
                "duration": 3.5,
                "type": "narration",
            },
            {
                "id": 5,
                "text": "ご不明な点がございましたら、いつでもご相談ください。ありがとうございました。",
                "duration": 4.0,
                "type": "narration",
            },
        ]
        
        return {
            "segments": segments,
            "total_duration": sum(s["duration"] for s in segments),
        }


def generate_script(
    pdf_data: Dict[str, Any], 
    video_data: Optional[Dict[str, Any]] = None
) -> Dict[str, Any]:
    """Convenience function for script generation"""
    try:
        generator = ScriptGenerator()
        return generator.generate(pdf_data, video_data)
    except (ValueError, ImportError) as e:
        print(f"   ⚠️ Cannot initialize ScriptGenerator: {e}")
        print("   Using fallback template script...")
        generator_fallback = ScriptGenerator.__new__(ScriptGenerator)
        return generator_fallback._generate_fallback_script(pdf_data)