"""Text generation interface using Google Gemini via google-genai.

This module provides a small, testable SDK surface for text generation.
It relies on the same Client used elsewhere in this project and exposes a
single convenience function `generate_text`.

Design goals:
- Pure function style: explicit inputs, structured outputs.
- No prints/logging here; callers control logging.
- No network retries here; callers can implement as needed.
"""

from __future__ import annotations

from pathlib import Path
from typing import Any, Optional, Union, Type
import json

import logging
from google import genai
from google.genai import types
import os
DEFAULT_TEXT_MODEL = "gemini-2.5-pro"

# Supported providers for text generation
TEXT_PROVIDERS = ("gemini", "codex")


_logger = logging.getLogger(__name__)
_logger.setLevel(logging.INFO)
if not _logger.handlers:
    handler = logging.StreamHandler()
    handler.setFormatter(logging.Formatter("%(levelname)s:%(name)s:%(message)s"))
    _logger.addHandler(handler)
from ._rate_limit import rate_limited


@rate_limited(max_calls=8, period=60.0, key="api_sdk.generate_text")
def generate_text(
    prompt: str,
    *,
    api_key: Optional[str] = None,
    model: Optional[str] = None,
    temperature: Optional[float] = None,
    max_output_tokens: Optional[int] = None,
    top_p: Optional[float] = None,
    top_k: Optional[int] = None,
    seed: Optional[int] = None,
    stop_sequences: Optional[list[str]] = None,
    response_mime_type: Optional[str] = None,
    response_schema: Optional[Union[dict[str, Any], type, types.Schema]] = None,
    response_json_schema: Optional[Any] = None,
    config: Optional[Union[types.GenerateContentConfig, types.GenerateContentConfigDict]] = None,
    extra_config: Optional[dict[str, Any]] = None,
    provider: str = "gemini",
) -> str:
    """Generate text via the selected provider.

    Args:
        prompt: User prompt (required)
        api_key: Optional API key (Gemini only; relies on google-genai default credentials if None)
        model: Model name (Gemini only; default: gemini-2.5-pro)
        temperature: Sampling temperature
        max_output_tokens: Maximum tokens in output
        top_p: Nucleus sampling parameter
        top_k: Top-k sampling parameter
        seed: Deterministic seed (if supported)
        stop_sequences: Optional stop sequences
        response_mime_type: Optional response mime type (e.g. 'application/json' for Structured Output)
        response_schema: Optional response schema (Structured Output). Can be a JSON schema dict or a python type.
        response_json_schema: Optional response JSON schema (Structured Output).
        config: Optional GenerateContentConfig (merged with explicit args)
        extra_config: Optional dict merged into config (advanced / forward-compatible)
        provider: 'gemini' (default) or 'codex'

    Returns:
        Generated text string.
    """
    if not prompt or not isinstance(prompt, str):
        raise ValueError("prompt must be a non-empty string")

    provider = (provider or "gemini").lower().strip()
    if provider not in TEXT_PROVIDERS:
        raise ValueError(f"Unsupported provider: {provider}. Supported: {', '.join(TEXT_PROVIDERS)}")

    if provider == "codex":
        if any(x is not None for x in (response_mime_type, response_schema, response_json_schema, config, extra_config)):
            raise ValueError("Structured output/config is only supported by provider='gemini'")
        return _generate_text_via_codex(prompt)
    else:
        return _generate_text_via_gemini(
            prompt,
            api_key=api_key,
            model=model,
            temperature=temperature,
            max_output_tokens=max_output_tokens,
            top_p=top_p,
            top_k=top_k,
            seed=seed,
            stop_sequences=stop_sequences,
            response_mime_type=response_mime_type,
            response_schema=response_schema,
            response_json_schema=response_json_schema,
            config=config,
            extra_config=extra_config,
        )


def _generate_text_via_codex(prompt: str) -> str:
    import shutil
    import subprocess
    import tempfile

    _logger.info("generate_text: using codex provider")
    codex_path = shutil.which("codex")
    if not codex_path:
        raise RuntimeError("Codex CLI not found on PATH. Please install with `npm i -g @openai/codex` or `brew install codex` and log in.")

    with tempfile.TemporaryDirectory() as td:
        out_file = Path(td) / "last_message.txt"
        cmd = [
            codex_path,
            "-a",
            "never",
            "exec",
            "--skip-git-repo-check",
            "--sandbox",
            "read-only",
            "--output-last-message",
            str(out_file),
            prompt,
        ]
        try:
            proc = subprocess.run(
                cmd,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True,
                timeout=180,
            )
        except subprocess.TimeoutExpired as e:
            raise RuntimeError("Codex text generation timed out") from e

        if proc.returncode != 0:
            raise RuntimeError(f"Codex failed (exit {proc.returncode}): {proc.stderr.strip() or proc.stdout.strip()}")

        if out_file.exists():
            return out_file.read_text(encoding="utf-8").strip()
        return (proc.stdout or "").strip()


def _schema_to_json_schema(schema: Any) -> dict[str, Any]:
    if isinstance(schema, dict):
        return schema
    model_json_schema = getattr(schema, "model_json_schema", None)
    if callable(model_json_schema):
        return model_json_schema()
    raise TypeError("schema must be a JSON schema dict or a Pydantic model class")


def _generate_text_structured_via_codex(prompt: str, *, schema: Any) -> str:
    import shutil
    import subprocess
    import tempfile

    _logger.info("generate_text: using codex provider (structured)")
    codex_path = shutil.which("codex")
    if not codex_path:
        raise RuntimeError("Codex CLI not found on PATH. Please install with `npm i -g @openai/codex` or `brew install codex` and log in.")

    schema_dict = _schema_to_json_schema(schema)

    with tempfile.TemporaryDirectory() as td:
        schema_path = Path(td) / "schema.json"
        out_file = Path(td) / "last_message.json"
        schema_path.write_text(json.dumps(schema_dict, ensure_ascii=False), encoding="utf-8")

        cmd = [
            codex_path,
            "-a",
            "never",
            "exec",
            "--skip-git-repo-check",
            "--sandbox",
            "read-only",
            "--output-schema",
            str(schema_path),
            "--output-last-message",
            str(out_file),
            prompt,
        ]
        try:
            proc = subprocess.run(
                cmd,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True,
                timeout=180,
            )
        except subprocess.TimeoutExpired as e:
            raise RuntimeError("Codex structured text generation timed out") from e

        if proc.returncode != 0:
            raise RuntimeError(f"Codex failed (exit {proc.returncode}): {proc.stderr.strip() or proc.stdout.strip()}")

        if out_file.exists():
            return out_file.read_text(encoding="utf-8").strip()
        return (proc.stdout or "").strip()


def _generate_text_via_gemini(
    prompt: str,
    *,
    api_key: Optional[str] = None,
    model: Optional[str] = None,
    temperature: Optional[float] = None,
    max_output_tokens: Optional[int] = None,
    top_p: Optional[float] = None,
    top_k: Optional[int] = None,
    seed: Optional[int] = None,
    stop_sequences: Optional[list[str]] = None,
    response_mime_type: Optional[str] = None,
    response_schema: Optional[Union[dict[str, Any], type, types.Schema]] = None,
    response_json_schema: Optional[Any] = None,
    config: Optional[Union[types.GenerateContentConfig, types.GenerateContentConfigDict]] = None,
    extra_config: Optional[dict[str, Any]] = None,
) -> str:
    _logger.info("generate_text: using gemini provider")
    env_key = api_key or os.environ.get("GEMINI_API_KEY")
    if not env_key:
        raise RuntimeError("GEMINI_API_KEY is required for Gemini text generation")
    client = genai.Client(api_key=env_key)

    cfg_data: dict[str, Any] = {}
    if config is not None:
        if isinstance(config, dict):
            cfg_data.update(config)
        else:
            cfg_data.update(config.model_dump(exclude_none=True))
    if extra_config:
        cfg_data.update(extra_config)

    if temperature is not None:
        cfg_data["temperature"] = float(temperature)
    if max_output_tokens is not None:
        cfg_data["max_output_tokens"] = int(max_output_tokens)
    if top_p is not None:
        cfg_data["top_p"] = float(top_p)
    if top_k is not None:
        cfg_data["top_k"] = int(top_k)
    if seed is not None:
        cfg_data["seed"] = int(seed)
    if stop_sequences is not None:
        cfg_data["stop_sequences"] = list(stop_sequences)
    if response_mime_type is not None:
        cfg_data["response_mime_type"] = response_mime_type
    if response_schema is not None:
        cfg_data["response_schema"] = response_schema
    if response_json_schema is not None:
        cfg_data["response_json_schema"] = response_json_schema

    cfg = types.GenerateContentConfig.model_validate(cfg_data) if cfg_data else types.GenerateContentConfig()
    use_model = (model or DEFAULT_TEXT_MODEL) or ""
    if isinstance(use_model, str) and "2.5-pro" in use_model and getattr(cfg, "thinking_config", None) is None:
        cfg.thinking_config = types.ThinkingConfig(include_thoughts=False, thinking_budget=128)

    response = client.models.generate_content(
        model=use_model or DEFAULT_TEXT_MODEL,
        contents=prompt,
        config=cfg,
    )

    text = getattr(response, "text", None) or ""
    if not text:
        parts = []
        for cand in getattr(response, "candidates", []) or []:
            content = getattr(cand, "content", None)
            for part in getattr(content, "parts", []) or []:
                t = getattr(part, "text", None)
                if t:
                    parts.append(t)
            if parts:
                break
        if parts:
            text = "".join(parts)
    return text


@rate_limited(max_calls=8, period=60.0, key="api_sdk.generate_text_structured")
def generate_text_structured(
    prompt: str,
    *,
    api_key: Optional[str] = None,
    model: Optional[str] = None,
    schema: Optional[Union[dict[str, Any], Type[Any]]] = None,
    response_mime_type: str = "application/json",
    config: Optional[Union[types.GenerateContentConfig, types.GenerateContentConfigDict]] = None,
    extra_config: Optional[dict[str, Any]] = None,
    provider: str = "gemini",
) -> Any:
    """Generate structured output via Gemini.

    Supported providers:
    - gemini: uses google-genai Structured Output (response_mime_type + response_schema/json_schema)
    - codex: uses Codex CLI Structured Output (--output-schema)

    Returns a dict (or a pydantic model instance when `schema` supports it).
    """
    if schema is None:
        raise ValueError("schema is required for structured output")

    provider = (provider or "gemini").lower().strip()
    if provider not in TEXT_PROVIDERS:
        raise ValueError(f"Unsupported provider: {provider}. Supported: {', '.join(TEXT_PROVIDERS)}")

    if provider == "codex":
        text = _generate_text_structured_via_codex(prompt, schema=schema)
    else:
        # If schema is a JSON schema dict, prefer response_json_schema. If it's a model/type, pass via response_schema.
        json_schema = schema if isinstance(schema, dict) else None
        typed_schema = None if isinstance(schema, dict) else schema
        text = generate_text(
            prompt,
            api_key=api_key,
            model=model,
            response_mime_type=response_mime_type,
            response_schema=typed_schema,
            response_json_schema=json_schema,
            config=config,
            extra_config=extra_config,
            provider="gemini",
        )

    # If schema looks like a pydantic model class, validate into it.
    model_validate_json = getattr(schema, "model_validate_json", None)
    if callable(model_validate_json):
        return model_validate_json(text)

    return json.loads(text)


__all__ = [
    "generate_text",
    "generate_text_structured",
    "DEFAULT_TEXT_MODEL",
    "TEXT_PROVIDERS",
]
