"""
LLM client abstraction for Lattice.

This module implements LLM completion via litellm or CLI passthrough:
- Thin wrapper around litellm.completion()
- CLI passthrough mode when config.model starts with "cli:"
- Provider encoded in config.model (e.g., "openrouter/google/gemini-3-flash-preview")
- API key read from environment variable
- Custom endpoints supported via config.base_url

CLI passthrough modes:
- "cli:claude" → claude -p <prompt> --output-format json --allowedTools ""
- "cli:opencode" → opencode run <prompt> --format json
- "cli:opencode:<model>" → opencode run <prompt> --model <model> --format json

All functions are Shell layer — they perform I/O and return Result[T, E].

Reference: RFC-002 §3.1 Configuration Schema, §4.3 Project Evolution
"""

from __future__ import annotations

import json
import os
import shutil
import subprocess
from typing import TYPE_CHECKING, Any

from returns.result import Failure, Result, Success

from lattice.core.model_utils import extract_provider, normalize_temperature
from lattice.shell.config import resolve_api_key_with_fallback

if TYPE_CHECKING:
    from lattice.core.types.config import CompilerConfig


# @shell_complexity: CLI subprocess orchestration with error handling
def _run_cli_command(command: list[str], cli_name: str) -> Result[str, str]:
    """Run a CLI command and return the output.

    Args:
        command: Command and arguments to run.
        cli_name: Name of the CLI tool (for error messages).

    Returns:
        Success(stdout) or Failure(error_message).
    """
    try:
        result = subprocess.run(
            command,
            capture_output=True,
            text=True,
            timeout=300,  # 5 minute timeout
        )

        if result.returncode != 0:
            return Failure(
                f"{cli_name} exited with code {result.returncode}: {result.stderr}"
            )

        return Success(result.stdout)

    except subprocess.TimeoutExpired:
        return Failure(f"{cli_name} command timed out after 300 seconds")
    except FileNotFoundError:
        return Failure(f"{cli_name} CLI not found. Install it or check PATH.")
    except Exception as e:
        return Failure(f"Failed to run {cli_name}: {e}")


# @shell_complexity: CLI passthrough with multiple model formats and error handling
def _cli_complete(config: "CompilerConfig", prompt: str) -> Result[str, str]:
    """Execute LLM completion via CLI passthrough.

    Handles models starting with "cli:" prefix:
    - "cli:claude" → claude CLI
    - "cli:opencode" → opencode CLI (default model)
    - "cli:opencode:<model>" → opencode CLI with specific model

    Args:
        config: CompilerConfig with cli: prefixed model.
        prompt: The user prompt/message to send.

    Returns:
        Success(response_text) or Failure(error_message).
    """
    model = config.model

    if model == "cli:claude":
        # Check claude CLI availability
        if not shutil.which("claude"):
            return Failure("claude CLI not found. Install claude CLI or check PATH.")

        command = [
            "claude",
            "-p",
            prompt,
            "--output-format",
            "json",
            "--allowedTools",
            "",
        ]
        result = _run_cli_command(command, "claude")
        if isinstance(result, Failure):
            return result

        try:
            data = json.loads(result.unwrap())
            # Claude JSON output has 'content' field with text
            if "content" in data:
                # content can be a list of content blocks
                content = data["content"]
                if isinstance(content, list):
                    # Extract text from content blocks
                    texts = []
                    for block in content:
                        if isinstance(block, dict) and block.get("type") == "text":
                            texts.append(block.get("text", ""))
                    return Success("".join(texts))
                return Success(str(content))
            # Fallback to 'result' field if content not available
            if "result" in data:
                return Success(data["result"])
            return Failure(f"Unexpected claude output format: {list(data.keys())}")
        except json.JSONDecodeError as e:
            return Failure(f"Failed to parse claude JSON output: {e}")

    elif model.startswith("cli:opencode"):
        # Check opencode CLI availability
        if not shutil.which("opencode"):
            return Failure(
                "opencode CLI not found. Install opencode CLI or check PATH."
            )

        # Parse model: "cli:opencode" or "cli:opencode:<model>"
        parts = model.split(":", 2)
        if len(parts) == 3 and parts[2]:
            # Specific model requested
            specific_model = parts[2]
            command = [
                "opencode",
                "run",
                prompt,
                "--model",
                specific_model,
                "--format",
                "json",
            ]
        else:
            # Default model
            command = [
                "opencode",
                "run",
                prompt,
                "--format",
                "json",
            ]

        result = _run_cli_command(command, "opencode")
        if isinstance(result, Failure):
            return result

        try:
            # OpenCode returns JSONL (one JSON per line)
            # We need to find the "text" type message
            output = result.unwrap()
            lines = output.strip().split("\n")

            for line in lines:
                if not line.strip():
                    continue
                try:
                    data = json.loads(line)
                    # OpenCode output format: {"type": "text", "part": {"text": "..."}}
                    if data.get("type") == "text":
                        part = data.get("part", {})
                        if "text" in part:
                            return Success(part["text"])
                    # Fallback formats
                    if "content" in data:
                        return Success(data["content"])
                    if "response" in data:
                        return Success(data["response"])
                except json.JSONDecodeError:
                    continue

            return Failure(
                f"Unexpected opencode output format. Lines: {len(lines)}, First line: {lines[0][:100] if lines else 'empty'}"
            )
        except Exception as e:
            return Failure(f"Failed to parse opencode output: {e}")

    else:
        return Failure(
            f"Unknown CLI model: {model}. Supported: cli:claude, cli:opencode, cli:opencode:<model>"
        )


# @shell_complexity: litellm API call with config handling + error types + CLI passthrough
def llm_complete(
    config: "CompilerConfig",
    prompt: str,
    temperature: float = 0.7,
    max_tokens: int | None = None,
) -> Result[str, str]:
    """Execute LLM completion via litellm or CLI passthrough.

    When config.model starts with "cli:", uses CLI passthrough:
    - "cli:claude" → claude CLI
    - "cli:opencode" → opencode CLI (default model)
    - "cli:opencode:<model>" → opencode CLI with specific model

    Otherwise uses litellm.completion():
    - OpenAI, OpenRouter, Anthropic, Ollama, and 100+ providers
    - Provider encoded in model string (e.g., "openrouter/google/gemini-3-flash-preview")
    - API key from environment variable (config.api_key_env)
    - Custom base URL for self-hosted or custom endpoints

    Args:
        config: CompilerConfig with model, api_key_env, base_url.
        prompt: The user prompt/message to send.
        temperature: Sampling temperature (default 0.7, ignored for CLI).
        max_tokens: Maximum tokens to generate (optional, ignored for CLI).

    Returns:
        Success(response_text) with the completion text.
        Failure(error_message) on API error.

    Example:
        >>> from lattice.core.types.config import CompilerConfig
        >>> # Via litellm
        >>> config = CompilerConfig(model="openrouter/google/gemini-3-flash-preview", api_key_env="OPENROUTER_API_KEY")
        >>> # result = llm_complete(config, "Hello, world!")  # Requires API key
        >>>
        >>> # Via CLI
        >>> cli_config = CompilerConfig(model="cli:claude", api_key_env="")  # No API key needed
        >>> # result = llm_complete(cli_config, "Hello, world!")  # Requires claude CLI
    """
    # Check for CLI passthrough mode
    if config.model.startswith("cli:"):
        return _cli_complete(config, prompt)

    try:
        import litellm

        # Resolve API key with full priority chain:
        # 1. config.api_key (with variable resolution)
        # 2. auth.json for provider
        # 3. config.api_key_env (env var)
        # 4. litellm default handling (None)
        provider = extract_provider(config.model)
        api_key_result = resolve_api_key_with_fallback(
            config_key=config.api_key,
            provider=provider,
            env_var=config.api_key_env if config.api_key_env else None,
        )
        if isinstance(api_key_result, Failure):
            return Failure(api_key_result.failure())
        api_key = api_key_result.unwrap()

        # Build messages
        messages: list[dict[str, str]] = [{"role": "user", "content": prompt}]

        # Normalize temperature for models with restrictions (e.g., GPT-5 requires temp=1)
        effective_temp = normalize_temperature(config.model, temperature)

        # Build kwargs for litellm
        kwargs: dict[str, Any] = {
            "model": config.model,
            "messages": messages,
            "temperature": effective_temp,
        }

        if api_key:
            kwargs["api_key"] = api_key
        if config.base_url:
            kwargs["api_base"] = config.base_url
        if max_tokens is not None:
            kwargs["max_tokens"] = max_tokens

        # OpenRouter special handling
        # OpenRouter requires OPENROUTER_API_KEY env var. When using a different
        # source (auth.json, non-standard env var), set it explicitly.
        if config.model.startswith("openrouter/") and api_key:
            os.environ["OPENROUTER_API_KEY"] = api_key

        # Ollama Cloud special handling
        # If using ollama/ prefix with OLLAMA_CLOUD_API_KEY, auto-set base_url
        if config.model.startswith("ollama/") and not config.base_url:
            if os.environ.get("OLLAMA_CLOUD_API_KEY"):
                # Note: litellm appends /api/ automatically for ollama models
                kwargs["api_base"] = "https://ollama.com"

        # Call litellm.completion
        response = litellm.completion(**kwargs)

        # Extract content from response
        content = response["choices"][0]["message"]["content"]

        return Success(content)

    except ImportError:
        return Failure("litellm not installed")
    except KeyError as e:
        return Failure(f"Unexpected response format: missing {e}")
    except Exception as e:
        return Failure(f"LLM completion error: {e}")


# @shell_complexity: litellm streaming + callback handling
def llm_complete_stream(
    config: "CompilerConfig",
    prompt: str,
    temperature: float = 0.7,
    max_tokens: int | None = None,
) -> Result[str, str]:
    """Execute LLM completion with streaming (collects all chunks).

    Uses litellm streaming mode. Collects all chunks and returns
    the complete response.

    Args:
        config: CompilerConfig with model, api_key, api_key_env, base_url.
        prompt: The user prompt/message to send.
        temperature: Sampling temperature (default 0.7).
        max_tokens: Maximum tokens to generate (optional).

    Returns:
        Success(response_text) with the complete completion.
        Failure(error_message) on API error.
    """
    try:
        import litellm

        # Resolve API key with full priority chain:
        # 1. config.api_key (with variable resolution)
        # 2. auth.json for provider
        # 3. config.api_key_env (env var)
        # 4. litellm default handling (None)
        provider = extract_provider(config.model)
        api_key_result = resolve_api_key_with_fallback(
            config_key=config.api_key,
            provider=provider,
            env_var=config.api_key_env if config.api_key_env else None,
        )
        if isinstance(api_key_result, Failure):
            return Failure(api_key_result.failure())
        api_key = api_key_result.unwrap()

        # Build messages
        messages: list[dict[str, str]] = [{"role": "user", "content": prompt}]

        # Normalize temperature for models with restrictions (e.g., GPT-5 requires temp=1)
        effective_temp = normalize_temperature(config.model, temperature)

        # Build kwargs for litellm
        kwargs: dict[str, Any] = {
            "model": config.model,
            "messages": messages,
            "temperature": effective_temp,
            "stream": True,
        }

        if api_key:
            kwargs["api_key"] = api_key
        if config.base_url:
            kwargs["api_base"] = config.base_url
        if max_tokens is not None:
            kwargs["max_tokens"] = max_tokens

        # Call litellm.completion with streaming
        response = litellm.completion(**kwargs)

        # Collect chunks
        chunks: list[str] = []
        for chunk in response:
            delta = chunk["choices"][0].get("delta", {})
            content = delta.get("content", "")
            if content:
                chunks.append(content)

        return Success("".join(chunks))

    except ImportError:
        return Failure("litellm not installed")
    except KeyError as e:
        return Failure(f"Unexpected response format: missing {e}")
    except Exception as e:
        return Failure(f"LLM completion error: {e}")


__all__ = [
    "llm_complete",
    "llm_complete_stream",
]
