"""
zai_judges.py - Multi-judge decision layer using Z.ai Anthropic-compatible API.

Used to rank palette candidates before generation so the system chooses a
coherent sonic direction instead of mixing unrelated local material.

Features:
- Exponential backoff retry for 429 errors (max 3 retries)
- Local cache with TTL for repeated prompts
- Clean heuristic fallback if API fails
"""

from __future__ import annotations

import hashlib
import json
import logging
import os
import re
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import Any, Dict, List, Optional, Tuple
from urllib.error import HTTPError, URLError
from urllib.request import Request, urlopen

logger = logging.getLogger("ZAIJudges")

# Cache configuration
CACHE_TTL_SECONDS = 600  # 10 minutes TTL for cache entries (increased from 5 min)
_cache: Dict[str, Tuple[Dict[str, Any], float]] = {}

# Retry configuration
MAX_RETRIES = 3
BACKOFF_DELAYS = [0.5, 1.0, 2.0]  # Reduced from [1.0, 2.0, 4.0] - cuts max backoff from 7s to 3.5s


def _generate_cache_key(system_prompt: str, user_payload: Dict[str, Any]) -> str:
    """Generate a deterministic cache key from prompt and payload."""
    # Include genre, style, candidates in key for palette context
    key_data = {
        "system": system_prompt[:200],  # First 200 chars of prompt
        "genre": user_payload.get("request", {}).get("genre", ""),
        "style": user_payload.get("request", {}).get("style", ""),
        "bpm": user_payload.get("request", {}).get("bpm", 0),
        "key": user_payload.get("request", {}).get("key", ""),
        "judge_role": user_payload.get("judge_role", ""),
        "candidate_ids": [
            c.get("id", "") 
            for c in user_payload.get("candidates", [])[:4]
        ],
    }
    key_string = json.dumps(key_data, sort_keys=True, ensure_ascii=True)
    return hashlib.sha256(key_string.encode("utf-8")).hexdigest()[:32]


def _get_cached_result(cache_key: str) -> Optional[Dict[str, Any]]:
    """Retrieve cached result if not expired."""
    if cache_key not in _cache:
        return None
    
    result, timestamp = _cache[cache_key]
    if time.time() - timestamp > CACHE_TTL_SECONDS:
        # Expired, remove from cache
        del _cache[cache_key]
        return None
    
    logger.debug("Cache hit for key: %s...", cache_key[:8])
    return result


def _set_cached_result(cache_key: str, result: Dict[str, Any]) -> None:
    """Store result in cache with current timestamp."""
    _cache[cache_key] = (result, time.time())
    logger.debug("Cached result for key: %s...", cache_key[:8])


def _resolve_messages_url() -> str:
    base = str(os.getenv("ANTHROPIC_BASE_URL", "https://api.z.ai/api/anthropic")).strip().rstrip("/")
    if base.endswith("/v1/messages"):
        return base
    if base.endswith("/v1"):
        return base + "/messages"
    return base + "/v1/messages"


def _extract_json_object(text: str) -> Dict[str, Any]:
    candidate = str(text or "").strip()
    if not candidate:
        return {}
    try:
        return json.loads(candidate)
    except Exception:
        pass

    match = re.search(r"\{.*\}", candidate, re.DOTALL)
    if not match:
        return {}
    try:
        return json.loads(match.group(0))
    except Exception:
        return {}


class ZAIJudgePanel:
    def __init__(self) -> None:
        self.base_url = _resolve_messages_url()
        self.auth_token = (
            os.getenv("ANTHROPIC_AUTH_TOKEN")
            or os.getenv("ZAI_API_KEY")
            or os.getenv("ANTHROPIC_API_KEY")
            or ""
        ).strip()
        self.model = str(os.getenv("ANTHROPIC_MODEL", "glm-5.1")).strip() or "glm-5.1"
        self.timeout = float(os.getenv("API_TIMEOUT_MS", "300000")) / 1000.0

    @property
    def available(self) -> bool:
        return bool(self.auth_token)

    def _call(self, system_prompt: str, user_payload: Dict[str, Any]) -> Dict[str, Any]:
        """Call Z.ai API with caching, retry logic, and exponential backoff for 429 errors."""
        if not self.available:
            return {}
        
        # Check cache first
        cache_key = _generate_cache_key(system_prompt, user_payload)
        cached = _get_cached_result(cache_key)
        if cached is not None:
            return cached

        body = {
            "model": self.model,
            "max_tokens": 550,
            "temperature": 0.2,
            "system": system_prompt,
            "messages": [
                {
                    "role": "user",
                    "content": json.dumps(user_payload, ensure_ascii=True),
                }
            ],
        }

        request = Request(
            self.base_url,
            data=json.dumps(body).encode("utf-8"),
            headers={
                "Content-Type": "application/json",
                "x-api-key": self.auth_token,
                "anthropic-version": "2023-06-01",
            },
            method="POST",
        )

        # Retry loop with exponential backoff
        last_error = None
        for attempt in range(MAX_RETRIES + 1):
            try:
                with urlopen(request, timeout=self.timeout) as response:
                    payload = json.loads(response.read().decode("utf-8", errors="replace"))
                
                text_chunks: List[str] = []
                for item in payload.get("content", []) or []:
                    if isinstance(item, dict) and item.get("type") == "text":
                        text_chunks.append(str(item.get("text", "")))
                
                result = _extract_json_object("\n".join(text_chunks))
                
                # Cache successful result
                if result:
                    _set_cached_result(cache_key, result)
                
                return result
                
            except HTTPError as error:
                last_error = error
                # Check for 429 (Too Many Requests)
                if hasattr(error, 'code') and error.code == 429:
                    if attempt < MAX_RETRIES:
                        delay = BACKOFF_DELAYS[min(attempt, len(BACKOFF_DELAYS) - 1)]
                        logger.warning("Judge API 429 on attempt %d/%d, retrying in %.1fs...", 
                                     attempt + 1, MAX_RETRIES, delay)
                        time.sleep(delay)
                        continue
                    else:
                        logger.error("Judge API 429 exhausted all %d retries", MAX_RETRIES)
                else:
                    logger.warning("Judge API HTTP error: %s (code: %s)", error, 
                                 getattr(error, 'code', 'unknown'))
                    break
                    
            except (URLError, TimeoutError) as error:
                last_error = error
                logger.warning("Judge API request failed: %s", error)
                break
                
            except Exception as error:
                last_error = error
                logger.warning("Judge API unexpected error: %s", error)
                break
        
        # All retries exhausted or non-retryable error
        logger.warning("Judge API call failed after %d attempt(s): %s", 
                     attempt + 1, last_error)
        return {}

    def judge_palette_candidates(
        self,
        genre: str,
        style: str,
        bpm: float,
        key: str,
        candidates: List[Dict[str, Any]],
        trend_context: Optional[Dict[str, Any]] = None,
    ) -> Dict[str, Any]:
        if not candidates:
            return {
                "available": False,
                "selected_candidate_id": "",
                "judges": [],
                "aggregate": {},
                "directives": {},
            }

        if not self.available:
            top = candidates[0]
            return {
                "available": False,
                "selected_candidate_id": top.get("id", ""),
                "judges": [],
                "aggregate": {
                    "selected_candidate_id": top.get("id", ""),
                    "score": float(top.get("score", 0.0)),
                    "mode": "heuristic_fallback",
                },
                "directives": {
                    "rhythm_density": "focused",
                    "bass_motion": "syncopated",
                    "arrangement_emphasis": ["intro", "build", "drop", "break", "drop", "outro"],
                    "vocal_strategy": "supportive",
                },
            }

        shortlist = candidates[:4]
        common_payload = {
            "request": {
                "genre": genre,
                "style": style,
                "bpm": bpm,
                "key": key,
            },
            "trend_context": trend_context or {},
            "candidates": shortlist,
            "response_contract": {
                "selected_candidate_id": "string",
                "score": "number_0_to_10",
                "strengths": ["string"],
                "weaknesses": ["string"],
                "directives": {
                    "rhythm_density": "string",
                    "bass_motion": "string",
                    "vocal_strategy": "string",
                    "arrangement_emphasis": ["string"],
                },
            },
        }

        judge_specs = [
            (
                "rhythm",
                (
                    "You are a reggaeton rhythm judge. Choose the palette candidate that will "
                    "produce the strongest dembow pocket, drum/bass chemistry and rhythmic coherence. "
                    "Respond as JSON only."
                ),
            ),
            (
                "harmony",
                (
                    "You are a reggaeton harmony and hook judge. Choose the palette candidate that will "
                    "produce the best tonal fit, melodic identity and vocal/music compatibility. "
                    "Respond as JSON only."
                ),
            ),
            (
                "arrangement",
                (
                    "You are a reggaeton arrangement judge. Choose the palette candidate that best supports "
                    "professional intro/build/drop/break/drop/outro pacing and section contrast. "
                    "Respond as JSON only."
                ),
            ),
        ]

        judge_results: List[Dict[str, Any]] = []
        with ThreadPoolExecutor(max_workers=min(3, len(judge_specs))) as executor:
            future_map = {
                executor.submit(self._call, prompt, {**common_payload, "judge_role": judge_name}): judge_name
                for judge_name, prompt in judge_specs
            }
            for future in as_completed(future_map):
                judge_name = future_map[future]
                try:
                    result = future.result() or {}
                except Exception as error:
                    logger.warning("Judge future failed (%s): %s", judge_name, error)
                    result = {}
                if result:
                    result["judge"] = judge_name
                    judge_results.append(result)

        if not judge_results:
            top = shortlist[0]
            return {
                "available": False,
                "selected_candidate_id": top.get("id", ""),
                "judges": [],
                "aggregate": {
                    "selected_candidate_id": top.get("id", ""),
                    "score": float(top.get("score", 0.0)),
                    "mode": "api_failed_heuristic_fallback",
                },
                "directives": {
                    "rhythm_density": "focused",
                    "bass_motion": "syncopated",
                    "arrangement_emphasis": ["intro", "build", "drop", "break", "drop", "outro"],
                    "vocal_strategy": "supportive",
                },
            }

        vote_counter: Dict[str, float] = {}
        directives: Dict[str, Any] = {}
        strengths: List[str] = []
        weaknesses: List[str] = []

        for result in judge_results:
            candidate_id = str(result.get("selected_candidate_id", "")).strip()
            score = float(result.get("score", 0.0) or 0.0)
            if candidate_id:
                vote_counter[candidate_id] = vote_counter.get(candidate_id, 0.0) + max(0.1, score)
            strengths.extend(str(item) for item in result.get("strengths", []) or [])
            weaknesses.extend(str(item) for item in result.get("weaknesses", []) or [])
            for key_name, value in dict(result.get("directives", {}) or {}).items():
                if key_name not in directives and value not in (None, "", []):
                    directives[key_name] = value

        selected_candidate_id = max(vote_counter.items(), key=lambda item: item[1])[0] if vote_counter else shortlist[0].get("id", "")
        aggregate_score = round(sum(float(result.get("score", 0.0) or 0.0) for result in judge_results) / len(judge_results), 2)

        return {
            "available": True,
            "model": self.model,
            "selected_candidate_id": selected_candidate_id,
            "judges": judge_results,
            "aggregate": {
                "selected_candidate_id": selected_candidate_id,
                "score": aggregate_score,
                "strengths": list(dict.fromkeys(strengths))[:10],
                "weaknesses": list(dict.fromkeys(weaknesses))[:10],
            },
            "directives": directives,
        }
