"""Gemini Flash function-calling backend for ORBIT Assist.

Uses the Gemini REST API directly (no SDK dependency) to perform
multi-turn function calling. This is the default backend since it
reuses the existing GOOGLE_API_KEY.

Handles known Gemini API quirks:
- Response key casing varies between API versions (camelCase vs snake_case)
- Large tool result payloads must be truncated to avoid context overflow
- Subprocess env must match the user's active Python environment
"""

from __future__ import annotations

import json
import logging
import urllib.request

from orbit.assist.tools import TOOL_DECLARATIONS, execute_tool

logger = logging.getLogger(__name__)

_GEMINI_URL = (
    "https://generativelanguage.googleapis.com/v1beta/"
    "models/{model}:generateContent?key={key}"
)

# Default model for function calling — Flash is fast and good at tool use
DEFAULT_MODEL = "gemini-2.5-flash"

# Maximum size (chars) of a single tool result before summarization kicks in.
# Gemini Flash has generous context but large payloads degrade quality and
# increase latency. 12K chars (~3K tokens) is a safe ceiling per tool result.
_MAX_TOOL_RESULT_CHARS = 12_000


class GeminiBackend:
    """Gemini Flash function-calling backend.

    Implements the agentic loop: send message -> model returns function_call ->
    execute tool -> send result back -> repeat until model gives final text.
    """

    def __init__(self, api_key: str, model: str = DEFAULT_MODEL):
        self.api_key = api_key
        self.model = model
        self.history: list[dict] = []

    def run_turn(
        self,
        user_message: str,
        system_prompt: str,
        on_tool_call: callable | None = None,
    ) -> str:
        """Run one turn of the agentic loop.

        May involve multiple tool calls before returning a final text response.

        Args:
            user_message: The user's input.
            system_prompt: Full system prompt with context.
            on_tool_call: Optional callback(tool_name, args, result) for UI updates.

        Returns:
            Final text response from the model.
        """
        self.history.append({
            "role": "user",
            "parts": [{"text": user_message}],
        })

        max_iterations = 15
        for iteration in range(max_iterations):
            response = self._call_gemini(system_prompt)
            if response is None:
                error_msg = "Failed to get a response from Gemini. Check your API key and network."
                self.history.append({"role": "model", "parts": [{"text": error_msg}]})
                return error_msg

            candidate = response.get("candidates", [{}])[0]
            content = candidate.get("content", {})
            parts = content.get("parts", [])

            if not parts:
                # Check for blocked / empty responses
                finish_reason = candidate.get("finishReason", "")
                if finish_reason == "SAFETY":
                    error_msg = "Response blocked by safety filter. Try rephrasing."
                else:
                    error_msg = "Empty response from Gemini."
                self.history.append({"role": "model", "parts": [{"text": error_msg}]})
                return error_msg

            # Extract function calls — handle BOTH key casing variants.
            # Gemini v1beta uses camelCase ("functionCall"), but some API
            # versions or SDK wrappers use snake_case ("function_call").
            function_calls = _extract_function_calls(parts)

            if not function_calls:
                # Final text response — no more tool calls
                text_parts = [
                    p.get("text", "")
                    for p in parts
                    if "text" in p
                ]
                final_text = "\n".join(text_parts).strip()
                self.history.append({"role": "model", "parts": parts})
                return final_text if final_text else "(No text in response)"

            # Model wants to call tools — add its response to history
            self.history.append({"role": "model", "parts": parts})

            # Execute each function call and collect results
            function_response_parts = []
            for tool_name, tool_args in function_calls:
                logger.debug("Tool call: %s(%s)", tool_name, tool_args)
                result = execute_tool(tool_name, tool_args)

                if on_tool_call:
                    on_tool_call(tool_name, tool_args, result)

                # Summarize large payloads to avoid context overflow
                result = _truncate_result(result, tool_name)

                function_response_parts.append({
                    "functionResponse": {
                        "name": tool_name,
                        "response": result,
                    }
                })

            # Send tool results back
            self.history.append({
                "role": "user",
                "parts": function_response_parts,
            })

        return "I've reached the maximum number of reasoning steps. Let me summarize what I've found so far."

    def _call_gemini(self, system_prompt: str) -> dict | None:
        """Make a single Gemini API call with function calling enabled."""
        url = _GEMINI_URL.format(model=self.model, key=self.api_key)

        body = {
            "system_instruction": {
                "parts": [{"text": system_prompt}],
            },
            "contents": self.history,
            "tools": [
                {"functionDeclarations": TOOL_DECLARATIONS},
            ],
            "generationConfig": {
                "temperature": 0.1,
                "maxOutputTokens": 4096,
            },
            "toolConfig": {
                "functionCallingConfig": {
                    "mode": "AUTO",
                },
            },
        }

        try:
            req = urllib.request.Request(
                url,
                data=json.dumps(body).encode(),
                headers={"Content-Type": "application/json"},
            )
            with urllib.request.urlopen(req, timeout=90) as resp:
                return json.loads(resp.read())
        except urllib.error.HTTPError as e:
            error_body = ""
            try:
                error_body = e.read().decode()
            except Exception:
                pass
            logger.debug("Gemini HTTP error %s: %s", e.code, error_body[:500])
            # Surface quota / auth errors to the user
            if e.code == 429:
                return _error_response("Rate limited by Gemini API. Wait a moment and try again.")
            if e.code in (401, 403):
                return _error_response("Gemini API key invalid or unauthorized. Check GOOGLE_API_KEY.")
            return None
        except Exception as e:
            logger.debug("Gemini call failed: %s", e)
            return None

    def reset(self) -> None:
        """Clear conversation history."""
        self.history.clear()


# ---------------------------------------------------------------------------
# Response parsing helpers
# ---------------------------------------------------------------------------


def _extract_function_calls(parts: list[dict]) -> list[tuple[str, dict]]:
    """Extract function calls from response parts.

    Handles both Gemini API key conventions:
      - camelCase: {"functionCall": {"name": ..., "args": ...}}
      - snake_case: {"function_call": {"name": ..., "args": ...}}
    """
    calls: list[tuple[str, dict]] = []
    for part in parts:
        fc = part.get("functionCall") or part.get("function_call")
        if fc and isinstance(fc, dict):
            name = fc.get("name", "")
            args = fc.get("args", {}) or {}
            # Gemini sometimes returns args as a JSON string instead of dict
            if isinstance(args, str):
                try:
                    args = json.loads(args)
                except (json.JSONDecodeError, ValueError):
                    args = {"raw_args": args}
            if name:
                calls.append((name, args))
    return calls


def _truncate_result(result: dict, tool_name: str) -> dict:
    """Truncate a tool result if it exceeds the safe payload size.

    Strategies by tool type:
    - orbit_analyze: keep grade/score/summary, truncate per-episode details
    - orbit_suggest: keep command + warnings, drop full hyperparameter tables
    - generic: serialize, check size, truncate deepest lists/strings
    """
    serialized = json.dumps(result, default=str)
    if len(serialized) <= _MAX_TOOL_RESULT_CHARS:
        return result

    logger.debug(
        "Truncating %s result from %d to ~%d chars",
        tool_name, len(serialized), _MAX_TOOL_RESULT_CHARS,
    )

    # Strategy 1: tool-specific key pruning
    if tool_name == "orbit_analyze":
        return _summarize_analyze_result(result)
    if tool_name == "scan_directory":
        return _summarize_scan_result(result)

    # Strategy 2: generic deep truncation
    return _deep_truncate(result)


def _summarize_analyze_result(result: dict) -> dict:
    """Keep the most important fields from a large analyze result."""
    keep_keys = {
        "grade", "score", "summary", "episode_count", "num_episodes",
        "num_frames", "fps", "robot_type", "action_dim", "state_dim",
        "cameras", "shapes", "consistency", "consistency_overall",
        "signal_health", "readiness", "predicted_success_range",
        "action_divergence", "recommendations", "warnings", "problems",
        "error", "output",
    }
    summary: dict = {}
    for k, v in result.items():
        if k in keep_keys:
            # Still truncate individual values if huge
            if isinstance(v, str) and len(v) > 2000:
                summary[k] = v[:2000] + "... (truncated)"
            elif isinstance(v, list) and len(v) > 20:
                summary[k] = v[:20]
                summary[f"{k}_note"] = f"Showing 20 of {len(v)} items"
            else:
                summary[k] = v
    if not summary:
        # Fallback: return first N chars as raw output
        return {"output": json.dumps(result, default=str)[:_MAX_TOOL_RESULT_CHARS]}
    return summary


def _summarize_scan_result(result: dict) -> dict:
    """Truncate scan results — cap dataset/output lists."""
    summary = dict(result)
    for key in ("datasets", "training_outputs", "config_files"):
        items = summary.get(key, [])
        if isinstance(items, list) and len(items) > 15:
            summary[key] = items[:15]
            summary[f"{key}_total"] = len(items)
    return summary


def _deep_truncate(result: dict, max_chars: int = _MAX_TOOL_RESULT_CHARS) -> dict:
    """Generic deep truncation: prune large strings and lists recursively."""

    def _shrink(obj, depth=0):
        if depth > 8:
            return "...(nested too deep)"
        if isinstance(obj, str):
            return obj[:1000] + "...(truncated)" if len(obj) > 1000 else obj
        if isinstance(obj, list):
            if len(obj) > 10:
                return [_shrink(item, depth + 1) for item in obj[:10]] + [
                    f"...({len(obj) - 10} more items)"
                ]
            return [_shrink(item, depth + 1) for item in obj]
        if isinstance(obj, dict):
            return {k: _shrink(v, depth + 1) for k, v in obj.items()}
        return obj

    shrunk = _shrink(result)
    # Verify it's actually smaller now
    if len(json.dumps(shrunk, default=str)) > max_chars * 1.5:
        # Nuclear option: just serialize and chop
        raw = json.dumps(result, default=str)[:max_chars]
        return {"output": raw, "_truncated": True}
    return shrunk


def _error_response(message: str) -> dict:
    """Build a fake Gemini response that surfaces an error as text."""
    return {
        "candidates": [{
            "content": {
                "parts": [{"text": message}],
            },
        }],
    }
