# Copyright (c) 2025, HaiyangLi <quantocean.li at gmail dot com>
# SPDX-License-Identifier: Apache-2.0

from __future__ import annotations

import asyncio  # Required for subprocess management (create_subprocess_exec)
import codecs
import contextlib
import inspect
import json  # Required for JSONDecoder.raw_decode() (streaming JSON parsing)
import logging
import shutil
from collections.abc import AsyncIterator, Callable
from dataclasses import (
    dataclass,
    field as datafield,
)
from pathlib import Path
from textwrap import shorten
from typing import Any, Literal

from pydantic import BaseModel, Field, model_validator

from lionpride.ln import json_dumps

HAS_GEMINI_CLI = False
GEMINI_CLI = None

if (g := (shutil.which("gemini") or "gemini")) and shutil.which(g):
    HAS_GEMINI_CLI = True
    GEMINI_CLI = g

logging.basicConfig(level=logging.INFO)
log = logging.getLogger("gemini-cli")

__all__ = (
    "GeminiChunk",
    "GeminiCodeRequest",
    "GeminiSession",
    "stream_gemini_cli",
)


class GeminiCodeRequest(BaseModel):
    """Request model for Gemini CLI execution."""

    # -- conversational bits -------------------------------------------------
    prompt: str = Field(description="The prompt for Gemini CLI")
    system_prompt: str | None = None

    # -- repo / workspace ----------------------------------------------------
    repo: Path = Field(default_factory=Path.cwd, exclude=True)
    ws: str | None = None  # sub-directory under repo
    include_directories: list[str] = Field(default_factory=list)

    # -- runtime & safety ----------------------------------------------------
    model: str | None = Field(
        default="gemini-2.5-pro",
        description="Gemini model to use (gemini-2.5-pro, gemini-2.5-flash, gemini-3-pro, etc.)",
    )
    yolo: bool = Field(
        default=False,
        description="Auto-approve all actions without confirmation (--yolo flag)",
    )
    approval_mode: Literal["suggest", "auto_edit", "full_auto"] | None = None
    debug: bool = False
    sandbox: bool = Field(
        default=True,
        description="Run in sandbox mode for safety",
    )

    # -- MCP integration -----------------------------------------------------
    mcp_tools: list[str] = Field(default_factory=list)

    # -- internal use --------------------------------------------------------
    verbose_output: bool = Field(default=False)
    cli_include_summary: bool = Field(default=False)

    @model_validator(mode="before")
    @classmethod
    def _validate_message_prompt(cls, data):
        """Convert messages format to prompt if needed."""
        if data.get("prompt"):
            return data

        if not (msg := data.get("messages")):
            raise ValueError("messages or prompt required")

        # Extract prompt from messages
        prompts = []
        for message in msg:
            if message["role"] != "system":
                content = message["content"]
                if isinstance(content, (dict, list)):
                    prompts.append(json_dumps(content))
                else:
                    prompts.append(content)
            elif message["role"] == "system" and not data.get("system_prompt"):
                data["system_prompt"] = message["content"]

        data["prompt"] = "\n".join(prompts)
        return data

    def cwd(self) -> Path:
        """Get working directory, validating workspace path."""
        if not self.ws:
            return self.repo

        ws_path = Path(self.ws)

        if ws_path.is_absolute():
            raise ValueError(f"Workspace path must be relative, got absolute: {self.ws}")

        if ".." in ws_path.parts:
            raise ValueError(f"Directory traversal detected in workspace path: {self.ws}")

        repo_resolved = self.repo.resolve()
        result = (self.repo / ws_path).resolve()

        try:
            result.relative_to(repo_resolved)
        except ValueError:
            raise ValueError(
                f"Workspace path escapes repository bounds. "
                f"Repository: {repo_resolved}, Workspace: {result}"
            )

        return result

    def as_cmd_args(self) -> list[str]:
        """Build argument list for the Gemini CLI."""
        args: list[str] = ["-p", self.prompt, "--output-format", "stream-json"]

        if self.model:
            args += ["-m", self.model]

        if self.yolo:
            args.append("--yolo")

        if self.approval_mode:
            args += ["--approval-mode", self.approval_mode]

        if self.debug:
            args.append("--debug")

        if not self.sandbox:
            args.append("--no-sandbox")

        for directory in self.include_directories:
            args += ["--include-directories", directory]

        return args


@dataclass
class GeminiChunk:
    """Low-level wrapper around every JSON object from the CLI."""

    raw: dict[str, Any]
    type: str
    # convenience views
    text: str | None = None
    tool_use: dict[str, Any] | None = None
    tool_result: dict[str, Any] | None = None


@dataclass
class GeminiSession:
    """Aggregated view of a whole CLI conversation."""

    session_id: str | None = None
    model: str | None = None

    # chronological log
    chunks: list[GeminiChunk] = datafield(default_factory=list)

    # materialized views
    messages: list[dict[str, Any]] = datafield(default_factory=list)
    tool_uses: list[dict[str, Any]] = datafield(default_factory=list)
    tool_results: list[dict[str, Any]] = datafield(default_factory=list)

    # final summary
    result: str = ""
    usage: dict[str, Any] = datafield(default_factory=dict)
    total_cost_usd: float | None = None
    num_turns: int | None = None
    duration_ms: int | None = None
    is_error: bool = False
    summary: dict | None = None

    def populate_summary(self) -> None:
        self.summary = _extract_summary(self)


def _extract_summary(session: GeminiSession) -> dict[str, Any]:
    """Extract summary from session data."""
    tool_counts = {}
    tool_details = []
    file_operations = {"reads": [], "writes": [], "edits": []}
    key_actions = []

    for tool_use in session.tool_uses:
        tool_name = tool_use.get("name", "unknown")
        tool_input = tool_use.get("input", {})
        tool_id = tool_use.get("id", "")

        tool_counts[tool_name] = tool_counts.get(tool_name, 0) + 1
        tool_details.append({"tool": tool_name, "id": tool_id, "input": tool_input})

        # Categorize by tool type
        if tool_name in ["read_file", "Read"]:
            file_path = tool_input.get("path", tool_input.get("file_path", "unknown"))
            file_operations["reads"].append(file_path)
            key_actions.append(f"Read {file_path}")

        elif tool_name in ["write_file", "Write"]:
            file_path = tool_input.get("path", tool_input.get("file_path", "unknown"))
            file_operations["writes"].append(file_path)
            key_actions.append(f"Wrote {file_path}")

        elif tool_name in ["edit_file", "Edit"]:
            file_path = tool_input.get("path", tool_input.get("file_path", "unknown"))
            file_operations["edits"].append(file_path)
            key_actions.append(f"Edited {file_path}")

        elif tool_name in ["run_shell_command", "shell", "Bash"]:
            command = tool_input.get("command", "")
            command_summary = command[:50] + "..." if len(command) > 50 else command
            key_actions.append(f"Ran: {command_summary}")

        elif tool_name.startswith("mcp_"):
            operation = tool_name.replace("mcp_", "")
            key_actions.append(f"MCP {operation}")

        else:
            key_actions.append(f"Used {tool_name}")

    key_actions = list(dict.fromkeys(key_actions)) if key_actions else ["No specific actions"]

    for op_type in file_operations:
        file_operations[op_type] = list(dict.fromkeys(file_operations[op_type]))

    result_summary = (session.result[:200] + "...") if len(session.result) > 200 else session.result

    return {
        "tool_counts": tool_counts,
        "tool_details": tool_details,
        "file_operations": file_operations,
        "key_actions": key_actions,
        "total_tool_calls": sum(tool_counts.values()),
        "result_summary": result_summary,
        "usage_stats": {
            "total_cost_usd": session.total_cost_usd,
            "num_turns": session.num_turns,
            "duration_ms": session.duration_ms,
            **session.usage,
        },
    }


async def _ndjson_from_cli(request: GeminiCodeRequest):
    """
    Yields each JSON object emitted by the Gemini CLI.

    Robust against UTF-8 splits and uses json.JSONDecoder.raw_decode.
    """
    workspace = request.cwd()
    workspace.mkdir(parents=True, exist_ok=True)

    proc = await asyncio.create_subprocess_exec(
        GEMINI_CLI,
        *request.as_cmd_args(),
        cwd=str(workspace),
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE,
    )

    decoder = codecs.getincrementaldecoder("utf-8")()
    json_decoder = json.JSONDecoder()
    buffer: str = ""

    try:
        while True:
            chunk = await proc.stdout.read(4096)
            if not chunk:
                break

            buffer += decoder.decode(chunk)

            while buffer:
                buffer = buffer.lstrip()
                if not buffer:
                    break
                try:
                    obj, idx = json_decoder.raw_decode(buffer)
                    yield obj
                    buffer = buffer[idx:]
                except json.JSONDecodeError:
                    break

        buffer += decoder.decode(b"", final=True)
        buffer = buffer.strip()
        if buffer:
            try:
                obj, idx = json_decoder.raw_decode(buffer)
                yield obj
            except json.JSONDecodeError:
                log.error("Skipped unrecoverable JSON tail: %.120s...", buffer)

        if await proc.wait() != 0:
            err = (await proc.stderr.read()).decode().strip()
            raise RuntimeError(err or "Gemini CLI exited non-zero")

    finally:
        with contextlib.suppress(ProcessLookupError):
            proc.terminate()
        await proc.wait()


async def stream_gemini_cli_events(request: GeminiCodeRequest):
    """Stream events from Gemini CLI."""
    if not GEMINI_CLI:
        raise RuntimeError("Gemini CLI not found (npm i -g @google/gemini-cli)")
    async for obj in _ndjson_from_cli(request):
        yield obj
    yield {"type": "done"}


async def _maybe_await(func, *args, **kw):
    """Call func which may be sync or async."""
    res = func(*args, **kw) if func else None
    if inspect.iscoroutine(res):
        await res


def _pp_text(text: str) -> None:
    print(f"\n> Gemini:\n{text}\n")


def _pp_tool_use(tu: dict[str, Any]) -> None:
    preview = shorten(str(tu.get("input", {})).replace("\n", " "), 130)
    print(f"- Tool Use - {tu.get('name', 'unknown')}: {preview}")


def _pp_tool_result(tr: dict[str, Any]) -> None:
    body_preview = shorten(str(tr.get("content", "")).replace("\n", " "), 130)
    status = "ERR" if tr.get("is_error") else "OK"
    print(f"- Tool Result - {status}: {body_preview}")


def _pp_final(sess: GeminiSession) -> None:
    usage = sess.usage or {}
    print(
        f"\n### Session complete\n"
        f"**Result:** {sess.result or ''}\n"
        f"- turns: {sess.num_turns}\n"
        f"- duration: {sess.duration_ms} ms\n"
        f"- tokens: {usage.get('input_tokens', 0)}/{usage.get('output_tokens', 0)}"
    )


async def stream_gemini_cli(
    request: GeminiCodeRequest,
    session: GeminiSession | None = None,
    *,
    on_text: Callable[[str], None] | None = None,
    on_tool_use: Callable[[dict[str, Any]], None] | None = None,
    on_tool_result: Callable[[dict[str, Any]], None] | None = None,
    on_final: Callable[[GeminiSession], None] | None = None,
) -> AsyncIterator[GeminiChunk | dict | GeminiSession]:
    """
    Consume the ND-JSON stream from Gemini CLI and return a populated GeminiSession.
    """
    if session is None:
        session = GeminiSession()

    stream = stream_gemini_cli_events(request)

    async for obj in stream:
        typ = obj.get("type", "unknown")
        chunk = GeminiChunk(raw=obj, type=typ)
        session.chunks.append(chunk)

        # Handle different event types based on Gemini CLI output format
        if typ in ("system", "init"):
            session.session_id = obj.get("session_id", obj.get("id"))
            session.model = obj.get("model")
            yield obj

        elif typ in ("message", "assistant"):
            msg = obj.get("message", obj)
            session.messages.append(msg)

            content = msg.get("content", "")
            if isinstance(content, str):
                chunk.text = content
                await _maybe_await(on_text, content)
                if request.verbose_output:
                    _pp_text(content)
            elif isinstance(content, list):
                for blk in content:
                    if isinstance(blk, dict):
                        btype = blk.get("type")
                        if btype == "text":
                            text = blk.get("text", "")
                            chunk.text = text
                            await _maybe_await(on_text, text)
                            if request.verbose_output:
                                _pp_text(text)
                        elif btype == "tool_use":
                            tu = {
                                "id": blk.get("id", ""),
                                "name": blk.get("name", ""),
                                "input": blk.get("input", {}),
                            }
                            chunk.tool_use = tu
                            session.tool_uses.append(tu)
                            await _maybe_await(on_tool_use, tu)
                            if request.verbose_output:
                                _pp_tool_use(tu)
            yield chunk

        elif typ in ("tool_call", "tool_use"):
            tu = {
                "id": obj.get("id", obj.get("tool_use_id", "")),
                "name": obj.get("name", obj.get("tool_name", "")),
                "input": obj.get("input", obj.get("args", {})),
            }
            chunk.tool_use = tu
            session.tool_uses.append(tu)
            await _maybe_await(on_tool_use, tu)
            if request.verbose_output:
                _pp_tool_use(tu)
            yield chunk

        elif typ == "tool_result":
            tr = {
                "tool_use_id": obj.get("tool_use_id", obj.get("id", "")),
                "content": obj.get("content", obj.get("result", "")),
                "is_error": obj.get("is_error", False),
            }
            chunk.tool_result = tr
            session.tool_results.append(tr)
            await _maybe_await(on_tool_result, tr)
            if request.verbose_output:
                _pp_tool_result(tr)
            yield chunk

        elif typ in ("result", "response"):
            session.result = obj.get("result", obj.get("response", "")).strip()
            session.usage = obj.get("usage", obj.get("stats", {}))
            session.total_cost_usd = obj.get("total_cost_usd", obj.get("cost"))
            session.num_turns = obj.get("num_turns", obj.get("turns"))
            session.duration_ms = obj.get("duration_ms", obj.get("duration"))
            session.is_error = obj.get("is_error", obj.get("error") is not None)

        elif typ == "error":
            session.is_error = True
            session.result = obj.get("message", obj.get("error", "Unknown error"))

        elif typ == "done":
            break

    await _maybe_await(on_final, session)
    if request.verbose_output:
        _pp_final(session)

    yield session
