"""
Infographic Title Generator - LLM Analyzer
Use LLM to analyze title structure and recommend suitable templates
"""

import json
import os
import random
import re
import requests
from typing import Dict, List, Optional


def _split_title_evenly(title: str, n_segments: int) -> List[str]:
    """按词把 title 大致均分成 n 段，作为最后兜底（语义切分也失败时使用）。"""
    if n_segments <= 0:
        n_segments = 1
    text = (title or "").strip()
    if not text:
        return [""] * n_segments
    words = re.split(r"\s+", text)
    if len(words) <= n_segments:
        out = list(words)
        while len(out) < n_segments:
            out.append("")
        return out
    base = len(words) // n_segments
    extra = len(words) % n_segments
    parts: List[str] = []
    idx = 0
    for i in range(n_segments):
        take = base + (1 if i < extra else 0)
        parts.append(" ".join(words[idx:idx + take]))
        idx += take
    return parts


# === Heuristic semantic splitter ===================================
# When LLM is disabled (default in the prod pipeline) we still want the
# title to break at natural language boundaries — colons, commas, the
# word "vs.", before prepositions / conjunctions — instead of a blind
# even-by-words split. This module-local DP picks N-1 break points to
# maximise a token-gap score table.

_PREPOSITIONS = frozenset({
    "of", "in", "by", "with", "for", "from", "on", "over", "per",
    "to", "at", "into", "across", "between", "after", "before",
    "during", "since", "through", "via", "under", "above",
})
_CONJUNCTIONS = frozenset({"and", "or", "but", "vs", "vs.", "versus", "&"})


def _gap_score(left_word: str, right_word: str) -> float:
    """Score the desirability of breaking after ``left_word``.

    Higher = more natural break point. Negative numbers mean "fairly
    arbitrary place to break" (we still let length-balance decide there).
    """
    lw = (left_word or "").rstrip()
    rw_lower = (right_word or "").lower().lstrip()
    if not lw or not rw_lower:
        return 0.0

    # Hard punctuation at the end of left word — strongest preference.
    last = lw[-1]
    if last == ":":
        return 1000.0
    if last == ";":
        return 900.0
    if last in {"—", "–"}:
        return 850.0
    if last == "?" or last == "!":
        return 800.0
    if last == "." and len(lw) > 2 and not lw.endswith(("Mr.", "Mrs.", "Dr.", "St.", "vs.")):
        # Likely sentence-ending period (skip common abbreviations and "vs.").
        return 700.0
    if last == ",":
        return 500.0

    # Soft preference: break BEFORE conjunctions / prepositions.
    rw_clean = rw_lower.rstrip(",.;:")
    if rw_clean in _CONJUNCTIONS:
        return 250.0
    if rw_clean in _PREPOSITIONS:
        return 150.0

    # Fallback: arbitrary mid-phrase break, no bonus, length balance wins.
    return 0.0


def _split_title_semantic(title: str, n_segments: int) -> List[str]:
    """Split ``title`` into ``n_segments`` segments by maximising a
    semantic-break score subject to length balance.

    Algorithm:
      1) Tokenise on whitespace; collect word widths in characters.
      2) For every gap i (between word[i] and word[i+1]) compute a
         positive bonus from punctuation/POS, then add a length-balance
         term that penalises uneven segment widths.
      3) Greedy DP: dp[k][i] = best score for splitting the first i words
         into k segments. Pick the top N-1 gaps that maximise the sum.

    Falls back to even split when title is too short for ``n_segments``.
    """
    if n_segments <= 0:
        n_segments = 1
    text = (title or "").strip()
    if not text:
        return [""] * n_segments
    words = re.split(r"\s+", text)
    W = len(words)
    if n_segments == 1:
        return [" ".join(words)]
    if W <= n_segments:
        out = list(words)
        while len(out) < n_segments:
            out.append("")
        return out

    # Per-word "width" in characters (proxy for visual width). Spaces between
    # words add 1 char.
    char_lens = [len(w) for w in words]
    # cumulative chars (each word + 1 space sep, except first), len = W+1
    cum = [0] * (W + 1)
    for i, c in enumerate(char_lens):
        cum[i + 1] = cum[i] + c + (1 if i > 0 else 0)
    total_chars = cum[W]
    target_seg = total_chars / n_segments

    def seg_chars(s: int, e: int) -> int:
        """Chars in words[s..e-1] joined by spaces."""
        if e <= s:
            return 0
        return cum[e] - cum[s] - (1 if s > 0 else 0)

    # Pre-compute gap bonuses for every internal gap i (break AFTER word i).
    gap_bonus = [0.0] * (W - 1)
    for i in range(W - 1):
        gap_bonus[i] = _gap_score(words[i], words[i + 1])

    def seg_score(s: int, e: int) -> float:
        """Score of segment words[s..e-1] viewed as one row."""
        if e <= s:
            return -1e9
        c = seg_chars(s, e)
        # Penalise distance from ideal segment length (per char).
        # Coefficient 5 keeps it on the same order as gap bonuses for typical
        # text (10-50 char titles), so a "+150 break before preposition"
        # never beats a 30-char balance issue.
        return -5.0 * abs(c - target_seg)

    # dp[k][i] = best total score using first i words in exactly k segments.
    # Track parent pointers to reconstruct the breaks.
    NEG_INF = float("-inf")
    dp = [[NEG_INF] * (W + 1) for _ in range(n_segments + 1)]
    parent = [[-1] * (W + 1) for _ in range(n_segments + 1)]
    dp[0][0] = 0.0
    for k in range(1, n_segments + 1):
        # Each segment must have at least 1 word; previous segments need
        # at least k-1 words; final segment must leave at least 1 word.
        for i in range(k, W + 1):
            best = NEG_INF
            best_j = -1
            for j in range(k - 1, i):
                if dp[k - 1][j] == NEG_INF:
                    continue
                # Break after word j (i.e. between word j-1 and word j when
                # j > 0); add the gap bonus for that break.
                gap = gap_bonus[j - 1] if j > 0 else 0.0
                cand = dp[k - 1][j] + seg_score(j, i) + gap
                if cand > best:
                    best = cand
                    best_j = j
            dp[k][i] = best
            parent[k][i] = best_j

    # Reconstruct boundaries.
    boundaries = []
    i = W
    for k in range(n_segments, 0, -1):
        j = parent[k][i]
        boundaries.append((j, i))
        i = j
    boundaries.reverse()

    parts = [" ".join(words[s:e]) for s, e in boundaries]
    # Safety net — should never trigger, but keep length contract.
    while len(parts) < n_segments:
        parts.append("")
    return parts[:n_segments]


class LLMAnalyzer:
    """LLM Title Analyzer"""
    
    def __init__(self, api_key=None, base_url=None, model=None):
        """
        Initialize LLM analyzer
        
        Args:
            api_key: API key (default uses built-in key)
            base_url: API base URL (default uses built-in URL)
            model: Model name (default uses gpt-5-mini)
        """
        # Use provided API configuration（环境变量优先，缺省禁用 LLM 走 fallback）
        self.api_key = api_key or os.environ.get('CHARTPIPELINE_LLM_API_KEY') or ''
        self.base_url = base_url or os.environ.get('CHARTPIPELINE_LLM_BASE_URL') or 'https://aihubmix.com/v1'
        self.model = model or os.environ.get('CHARTPIPELINE_LLM_MODEL') or 'deepseek-v3.2'
        self.llm_disabled = os.environ.get('CHARTPIPELINE_DISABLE_LLM', '1') == '1' or not self.api_key
        # 仅第一次失败时打印 fallback 提示，避免刷屏
        self._fallback_logged = False

    def _log_fallback_once(self, reason: str):
        if not self._fallback_logged:
            print(f"  ⚠️  LLM unavailable ({reason}), using local heuristic fallback for title styling")
            self._fallback_logged = True

    def _fallback_single(self, title: str, templates_info: Optional[List[Dict]]) -> Dict:
        """LLM 失败时的单模板 fallback。"""
        self._log_fallback_once("single-template")
        if templates_info:
            chosen = random.choice(templates_info)
            name = chosen['name']
            seg = int(chosen.get('segment_count') or len(chosen.get('segments', [])) or 1)
        else:
            name = 'infographic_standard'
            seg = 1
        return {
            'recommended_template': name,
            'title_split': _split_title_semantic(title, seg),
            'split_method': 'semantic_dp',
            'reasoning': 'fallback: LLM disabled, picked template by random sampling and split title via semantic DP',
        }

    def _fallback_top_k(self, title: str, all_templates_info: list, top_k: int) -> Dict:
        """LLM 失败时的 top-k fallback。"""
        self._log_fallback_once("top-k")
        if not all_templates_info:
            return {'recommendations': []}
        pool = list(all_templates_info)
        random.shuffle(pool)
        recs = []
        for info in pool[:max(1, top_k)]:
            seg = int(info.get('segment_count') or 1)
            recs.append({
                'template_name': info['name'],
                'title_split': _split_title_semantic(title, seg),
                'split_method': 'semantic_dp',
                'paraphrased_title': None,
                'confidence': 0.5,
                'reasoning': 'fallback: LLM disabled, random template + semantic DP split',
            })
        return {'recommendations': recs}
    
    def recommend_template(self, title: str, has_description: bool = True, templates_info: Optional[List[Dict]] = None) -> Dict:
        """
        Analyze title and recommend suitable template with splitting scheme
        
        Args:
            title: Main title text
            has_description: Whether description is provided
            templates_info: List of template info dicts (optional, for filtering by style)
        
        Returns:
            Dictionary containing:
            {
                'recommended_template': 'template_name',
                'title_split': ['line1', 'line2', ...],
                'reasoning': 'explanation'
            }
        """
        return self._analyze_with_llm(title, has_description, templates_info)
    
    def recommend_top_k_templates(self, title: str, all_templates_info: list, top_k: int = 3) -> Dict:
        """
        Analyze title and recommend top-k suitable templates with splitting schemes
        
        Args:
            title: Main title text
            all_templates_info: List of template info dicts with name, description, segment_roles, segments
            top_k: Number of templates to recommend (default 3)
        
        Returns:
            Dictionary containing:
            {
                'recommendations': [
                    {
                        'template_name': 'template1',
                        'title_split': ['line1', 'line2', ...],
                        'paraphrased_title': 'Restructured title or null',
                        'confidence': 0.95,
                        'reasoning': 'explanation'
                    },
                    ...
                ]
            }
        """
        return self._analyze_with_llm_top_k(title, all_templates_info, top_k)
    
    
    def _analyze_with_llm(self, title: str, has_description: bool, templates_info: Optional[List[Dict]] = None) -> Dict:
        """
        Use LLM API to analyze and recommend template
        """
        if self.llm_disabled:
            return self._fallback_single(title, templates_info)

        # Build prompt
        prompt = self._build_recommendation_prompt(title, has_description, templates_info)
        
        # Call LLM API
        try:
            response = self._query_llm(prompt)
            
            if response:
                # Parse JSON response
                try:
                    # Try to clean possible markdown code blocks
                    cleaned_response = response.strip()
                    if cleaned_response.startswith('```'):
                        lines = cleaned_response.split('\n')
                        cleaned_response = '\n'.join(lines[1:-1] if lines[-1].strip() == '```' else lines[1:])
                        cleaned_response = cleaned_response.replace('```json', '').replace('```', '').strip()
                    
                    result = json.loads(cleaned_response)
                    
                    # Validate return format
                    if 'recommended_template' in result and 'title_split' in result:
                        # Validate template name if templates_info is provided
                        if templates_info:
                            available_names = [t['name'] for t in templates_info]
                            if result['recommended_template'] not in available_names:
                                print(f"❌ Unknown template '{result['recommended_template']}'")
                                return None
                        return result
                    else:
                        print("❌ LLM response missing required fields")
                        return self._fallback_single(title, templates_info)
                        
                except json.JSONDecodeError as e:
                    print(f"❌ LLM response is not valid JSON")
                    print(f"Response: {response[:200]}...")
                    return self._fallback_single(title, templates_info)
            else:
                return self._fallback_single(title, templates_info)
                
        except Exception as e:
            print(f"❌ LLM analysis error: {e}")
            return self._fallback_single(title, templates_info)
    
    def _query_llm(self, prompt: str) -> Optional[str]:
        """
        Query LLM API
        
        Args:
            prompt: Prompt to send to LLM
        
        Returns:
            str: LLM response content
        """
        headers = {
            'Authorization': f'Bearer {self.api_key}',
            'Content-Type': 'application/json'
        }
        
        data = {
            'model': self.model,
            'messages': [
                {
                    'role': 'system', 
                    'content': 'You are a professional infographic design assistant, skilled at analyzing title structure and visual hierarchy. Always return valid JSON format only, without any markdown formatting or extra text.'
                },
                {
                    'role': 'user', 
                    'content': prompt
                }
            ],
            'temperature': 0.3
        }
        
        try:
            response = requests.post(
                f'{self.base_url}/chat/completions',
                headers=headers,
                json=data,
                timeout=30
            )
            response.raise_for_status()
            
            result = response.json()
            return result['choices'][0]['message']['content'].strip()
            
        except requests.exceptions.Timeout:
            print("❌ LLM API timeout")
            return None
        except requests.exceptions.HTTPError as e:
            print(f"❌ LLM API HTTP error: {e}")
            if hasattr(e.response, 'text'):
                print(f"   Response: {e.response.text[:200]}")
            return None
        except requests.exceptions.RequestException as e:
            print(f"❌ LLM API request error: {e}")
            return None
        except KeyError as e:
            print(f"❌ LLM API response format error: {e}")
            return None
    
    def _analyze_with_llm_top_k(self, title: str, all_templates_info: list, top_k: int = 3) -> Dict:
        """
        Use LLM API to recommend top-k templates
        
        Args:
            title: Main title text
            all_templates_info: List of template info dicts
            top_k: Number of templates to recommend
        
        Returns:
            Dictionary with recommendations list
        """
        if self.llm_disabled:
            return self._fallback_top_k(title, all_templates_info, top_k)

        # Build prompt for top-k recommendation
        prompt = self._build_top_k_prompt(title, all_templates_info, top_k)
        
        # Call LLM API
        try:
            response = self._query_llm(prompt)
            
            if response:
                # Parse JSON response
                try:
                    cleaned_response = response.strip()
                    if cleaned_response.startswith('```'):
                        lines = cleaned_response.split('\n')
                        cleaned_response = '\n'.join(lines[1:-1] if lines[-1].strip() == '```' else lines[1:])
                        cleaned_response = cleaned_response.replace('```json', '').replace('```', '').strip()
                    
                    result = json.loads(cleaned_response)
                    
                    # Validate return format
                    if 'recommendations' in result and isinstance(result['recommendations'], list):
                        return result
                    else:
                        print("❌ LLM response missing required 'recommendations' field")
                        return self._fallback_top_k(title, all_templates_info, top_k)
                        
                except json.JSONDecodeError as e:
                    print(f"❌ LLM response is not valid JSON")
                    return self._fallback_top_k(title, all_templates_info, top_k)
            else:
                return self._fallback_top_k(title, all_templates_info, top_k)
                
        except Exception as e:
            print(f"❌ LLM analysis error: {e}")
            return self._fallback_top_k(title, all_templates_info, top_k)
    
    def _build_recommendation_prompt(self, title: str, has_description: bool, templates_info: Optional[List[Dict]] = None) -> str:
        """Build LLM recommendation prompt"""
        
        # Build default templates info if not provided
        if not templates_info:
            # Use a basic set of templates
            templates_list = [
                {"name": "infographic_standard", "description": "Single line bold title", "segments": 1},
                {"name": "two_line_hierarchy", "description": "Two-line: small prefix + large main title", "segments": 2},
                {"name": "three_line_emphasis", "description": "Three-line: prefix + emphasized middle + suffix", "segments": 3}
            ]
        else:
            templates_list = templates_info
        
        templates_desc = "\n".join([
            f"- **{t['name']}**: {t.get('description', 'No description')}, segments: {t.get('segment_count', len(t.get('segments', [])))}"
            for t in templates_list
        ])
        
        prompt = f"""
Analyze this infographic title to recommend the layout template and split strategy.

Input Data:
- Title: "{title}"
- Has Subtitle/Description: {has_description}

### Task
1. Identify the **PRIMARY** focus (The "Hero" of the title) and **SECONDARY** context.
2. Recommend the best template (`infographic_standard`, `two_line_hierarchy`, or `three_line_emphasis`).
3. Split the title text strictly according to the template rules.

### 1. Semantic Analysis Rules
**PRIMARY Segment (The "Hero"):**
- The core metric, main topic, or "What" the chart shows.
- *Examples:* "GDP Growth", "Carbon Emissions", "Mobile Phones", "Annual Report".
- **Constraint:** NEVER split a noun phrase inside the Primary segment.

**SECONDARY Segment (The "Context"):**
- Modifiers, Prepositions, Locations, Time, or "Which/Where/When".
- *Examples:* "The United States of...", "Trends in...", "...of All Time", "Global...", "Top 10...".
- *Note:* If a title contains both a Metric (e.g., Inflation) and a Location (e.g., in Europe), treat the Metric as PRIMARY and Location as SECONDARY.

### 2. Template Selection Guidelines
**Option A: infographic_standard**
- *Criteria:* Short titles (1-5 words) OR titles that are a single cohesive phrase.
- *Split Strategy:* Return as a single string in the array.
- *Example:* ["Annual Report 2023"]

**Option B: two_line_hierarchy**
- *Criteria:* Strong prefix (Context) + Main Topic (Hero). Good for "The X of Y" structures.
- *Split Strategy:* - Line 1: Secondary Context (Prefix)
  - Line 2: Primary Topic (Hero)
- *Example:* ["The United States of", "Food Inflation"]

**Option C: three_line_emphasis**
- *Criteria:* Long titles (6+ words) with a "Sandwich" structure: Context + Hero + Context.
- *Split Strategy:*
  - Line 1: Secondary Context (Prefix)
  - Line 2: Primary Topic (Hero - Centerpiece)
  - Line 3: Secondary Context (Suffix)
- *Example:* ["The Best Selling", "Mobile Phones", "of All Time"]

### Output Format
Return ONLY valid JSON:
```json
{{
  "recommended_template": "template_name",
  "title_split": ["string1", "string2", ...],
  "reasoning": "Explain why this template was chosen based on the Primary/Secondary analysis."
}}
"""
        
        return prompt
    
    def _build_top_k_prompt(self, title: str, all_templates_info: list, top_k: int) -> str:
        """
        Build prompt for top-k template recommendation
        
        Args:
            title: Main title text
            all_templates_info: List of template info with segment_roles
            top_k: Number of recommendations
        
        Returns:
            Prompt string
        """
        # Build templates list for prompt
        templates_desc = []
        for info in all_templates_info:
            name = info['name']
            segment_roles = info.get('segment_roles', 'No role description available')
            segment_count = info.get('segment_count', 1)
            templates_desc.append(f"- **{name}** ({segment_count} segments): {segment_roles}")
        
        templates_text = "\n".join(templates_desc)
        
        prompt = f"""You are an expert in infographic title design and linguistic analysis.

**Task**: Analyze the given title and recommend the top {top_k} most suitable templates based on grammatical structure and semantic roles.

**Input Title**: "{title}"

**Available Templates** (with segment role descriptions):
{templates_text}

**Your Responsibilities**:

1. **Grammatical Analysis**: Parse the title to identify its grammatical components (subject, modifier, qualifier, possessive, prepositional phrase, etc.)

2. **Structural Matching**: Match the title's structure against each template's segment_roles based on:
   - Number of natural semantic units in the title
   - Grammatical roles of each unit (modifier, subject, qualifier, etc.)
   - Visual hierarchy implied by the content

3. **Title Adaptation**: If necessary, **paraphrase or restructure** the title to better fit the template:
   - Preserve the core meaning
   - Adjust phrasing to match segment roles
   - Optimize for visual impact

4. **Segmentation**: For each recommended template, split the title (original or paraphrased) into segments that match the template's segment_roles.

5. **Ranking**: Order recommendations by confidence, considering both grammatical fit and visual effectiveness.

**Output Format** (return ONLY valid JSON, no markdown):
{{
  "recommendations": [
    {{
      "template_name": "template_name",
      "title_split": ["segment1", "segment2", ...],
      "paraphrased_title": "Restructured title if modified, or null if unchanged",
      "confidence": 0.95,
      "reasoning": "Brief explanation of grammatical structure match and any paraphrasing done"
    }}
  ]
}}

**Critical Guidelines on Segment Importance**:
- **PRIMARY importance**: Core subject nouns, key entities, or concrete data/metrics that form the main topic
  * Examples: "Life Expectancy", "Mobile Phones", "Carbon Emissions", "Europe", "GDP Growth"
  * These are the focal subjects - the "what" the infographic is about
  * Will be rendered with larger, bolder, more prominent styling
  * **KEEP PRIMARY INTACT**: Do not split core noun phrases (e.g., keep "Life Expectancy" together)
  
- **SECONDARY importance**: Modifiers, time references, actions, prepositions, and contextual phrases
  * Examples: 
    - Time references: "Since 1950", "in 2023", "by 2030"
    - Actions/verbs: "Gains", "Changes", "Growth"
    - Prepositions: "by Country", "via Technology", "of All Time"
    - Quantifiers/modifiers: "Top 10", "Best-Selling", "Average"
  * These provide context but are NOT the main subject
  * Will be rendered with smaller, less prominent styling

**Key Principle**: Identify the SUBJECT first (PRIMARY), then treat everything else as decoration (SECONDARY).
  - Example: "Average Life Expectancy Gains Since 1950" → PRIMARY: "Life Expectancy", SECONDARY: "Average...Gains Since 1950"
  - Example: "Europe's Best-Selling Car Brands" → PRIMARY: "Car Brands" (or "Europe"), SECONDARY: "Best-Selling"

**Other Guidelines**:
- Ensure title_split length matches the template's segment count
- Consider grammatical roles (possessive, modifier, subject, etc.)
- Paraphrase when it improves fit without changing meaning
- Prioritize templates with natural grammatical alignment
- Provide diverse options (avoid similar templates)
- Order by confidence (highest first)
- **Identify topic keywords and metrics for PRIMARY segments, context for SECONDARY**
"""
        return prompt
    


# Convenience function