"""
gemini_cli_wrapper.py — Class wrapper gọi Gemini CLI từ Python.

Dùng cho:
  1. Lead Flow: thay thế LangChain model → gọi Gemini CLI trực tiếp
  2. Test Runner: sinh câu hỏi + đánh giá response

Yêu cầu: `gemini` CLI đã cài và có API key.
"""

import asyncio
import json
import logging
import re
import shutil
import subprocess
import tempfile

logger = logging.getLogger(__name__)


class GeminiCLI:
    """
    Wrapper class gọi Gemini CLI command.
    
    Có thể cắm vào bất kỳ đâu cần AI trả lời:
      - Lead Flow (thay LangChain)
      - Test Runner (sinh câu hỏi + đánh giá)
      - Bất kỳ module nào cần LLM
    
    Usage:
        gemini = GeminiCLI()
        
        # Text response
        answer = await gemini.ask("Tư vấn áo đi biển")
        
        # JSON response  
        data = await gemini.ask_json("Parse câu hỏi thành JSON: ...")
        
        # Sync version
        answer = gemini.ask_sync("Hello")
    """

    def __init__(self, model: str = "gemini-2.0-flash", timeout: int = 120):
        self.model = model
        self.timeout = timeout
        self._cli_path: str | None = None

    @property
    def cli_path(self) -> str:
        """Lazy load đường dẫn gemini CLI."""
        if self._cli_path is None:
            self._cli_path = shutil.which("gemini")
            if not self._cli_path:
                raise FileNotFoundError(
                    "Gemini CLI không tìm thấy trên PATH. "
                    "Cài đặt: npm install -g @anthropic-ai/gemini-cli"
                )
            logger.info("[GeminiCLI] Found: %s", self._cli_path)
        return self._cli_path

    # ── ASYNC ──

    async def ask(self, prompt: str) -> str:
        """Gọi Gemini CLI async, trả về text response."""
        env = {
            **__import__("os").environ,
            "NO_COLOR": "1",
            "TERM": "dumb",
            "FORCE_COLOR": "0",
        }
        # Pass via stdin to avoid "Command line is too long" on Windows
        cwd = tempfile.gettempdir()
        cmd = [self.cli_path]
        proc = await asyncio.create_subprocess_exec(
            *cmd,
            stdin=asyncio.subprocess.PIPE,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
            env=env,
            cwd=cwd,
        )
        stdout, stderr = await asyncio.wait_for(
            proc.communicate(input=prompt.encode("utf-8")), timeout=self.timeout
        )
        output = stdout.decode("utf-8", errors="replace").strip()
        err = stderr.decode("utf-8", errors="replace").strip()

        # Gemini CLI có thể exit non-zero nhưng vẫn có output (do warning)
        if output:
            if err:
                logger.debug("[GeminiCLI] stderr (ignored): %s", err[:200])
            return output

        if proc.returncode != 0:
            logger.error("[GeminiCLI] Error (code=%d): %s", proc.returncode, err)
            raise RuntimeError(f"Gemini CLI failed: {err}")

        logger.debug("[GeminiCLI] Response: %.100s...", output)
        return output

    async def ask_json(self, prompt: str) -> dict:
        """Gọi Gemini CLI, tự động parse response thành JSON."""
        raw = await self.ask(prompt)
        return self._parse_json(raw)

    # ── SYNC ──

    def ask_sync(self, prompt: str) -> str:
        """Phiên bản SYNC — subprocess.run trực tiếp."""
        cwd = tempfile.gettempdir()
        result = subprocess.run(
            [self.cli_path],
            input=prompt,
            capture_output=True,
            text=True,
            timeout=self.timeout,
            cwd=cwd,
        )
        if result.returncode != 0:
            raise RuntimeError(f"Gemini CLI failed: {result.stderr}")
        return result.stdout.strip()

    def ask_json_sync(self, prompt: str) -> dict:
        """Phiên bản SYNC — trả về JSON."""
        raw = self.ask_sync(prompt)
        return self._parse_json(raw)

    # ── JSON Parser ──

    @staticmethod
    def _parse_json(raw: str) -> dict:
        """
        Tìm và parse JSON từ output Gemini CLI.
        Hỗ trợ: ```json ... ```, raw JSON, hoặc { ... } trong text.
        """
        # 1. Tìm fenced code block
        json_match = re.search(r"```(?:json)?\s*\n([\s\S]*?)\n```", raw)
        if json_match:
            try:
                return json.loads(json_match.group(1).strip())
            except json.JSONDecodeError:
                pass

        # 2. Thử parse trực tiếp
        try:
            return json.loads(raw)
        except json.JSONDecodeError:
            pass

        # 3. Tìm { ... } trong text
        brace_match = re.search(r"\{[\s\S]*\}", raw)
        if brace_match:
            try:
                return json.loads(brace_match.group())
            except json.JSONDecodeError:
                pass

        logger.warning("[GeminiCLI] Cannot parse JSON: %.200s", raw)
        return {"raw_text": raw}


# ═══════════════════════════════════════════════
# LangChain BaseChatModel Adapter
# ═══════════════════════════════════════════════

from typing import Any, Optional

from langchain_core.language_models import BaseChatModel
from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, SystemMessage
from langchain_core.outputs import ChatGeneration, ChatResult
from langchain_core.runnables import RunnableLambda
from pydantic import PrivateAttr


class ChatGeminiCLI(BaseChatModel):
    """
    LangChain BaseChatModel adapter — bọc Gemini CLI subprocess.

    Dùng OAuth authentication từ Gemini CLI (KHÔNG cần GOOGLE_API_KEY).
    Hỗ trợ with_structured_output() qua JSON-in-prompt.

    Usage:
        llm = ChatGeminiCLI(model="gemini-2.0-flash")
        result = await llm.ainvoke([HumanMessage(content="hello")])

        # Structured output (thay thế function_calling)
        structured = llm.with_structured_output(MyPydanticModel)
        obj = await structured.ainvoke(messages)
    """

    model: str = "gemini-2.0-flash"
    timeout: int = 120
    temperature: float = 0
    max_output_tokens: int = 1500
    streaming: bool = False  # CLI không hỗ trợ streaming, chấp nhận param cho compat

    _cli: GeminiCLI = PrivateAttr()

    model_config = {"arbitrary_types_allowed": True}

    def __init__(self, **kwargs):
        # Compat: llm_factory truyền 'model' hoặc 'model_name'
        if "model_name" in kwargs and "model" not in kwargs:
            kwargs["model"] = kwargs.pop("model_name")
        elif "model_name" in kwargs:
            kwargs.pop("model_name")

        # Loại bỏ params không liên quan đến CLI
        for drop_key in ("google_api_key", "response_mime_type", "api_key", "model_kwargs"):
            kwargs.pop(drop_key, None)

        super().__init__(**kwargs)
        self._cli = GeminiCLI(model=self.model, timeout=self.timeout)
        logger.info("✅ [ChatGeminiCLI] Initialized | model=%s | OAuth via Gemini CLI", self.model)

    @property
    def _llm_type(self) -> str:
        return "gemini-cli"

    def _format_messages(self, messages: list[BaseMessage]) -> str:
        """Flatten LangChain messages thành 1 prompt string cho CLI."""
        parts = []
        for msg in messages:
            if isinstance(msg, SystemMessage):
                parts.append(f"[System Instructions]\n{msg.content}")
            elif isinstance(msg, HumanMessage):
                parts.append(f"[User]\n{msg.content}")
            elif isinstance(msg, AIMessage):
                parts.append(f"[Assistant]\n{msg.content}")
            else:
                parts.append(f"[{msg.type}]\n{msg.content}")
        return "\n\n".join(parts)

    def _generate(
        self,
        messages: list[BaseMessage],
        stop: Optional[list[str]] = None,
        run_manager: Any = None,
        **kwargs,
    ) -> ChatResult:
        """Sync generation via Gemini CLI subprocess."""
        prompt = self._format_messages(messages)
        response = self._cli.ask_sync(prompt)
        return ChatResult(
            generations=[ChatGeneration(message=AIMessage(content=response))]
        )

    async def _agenerate(
        self,
        messages: list[BaseMessage],
        stop: Optional[list[str]] = None,
        run_manager: Any = None,
        **kwargs,
    ) -> ChatResult:
        """Async generation via Gemini CLI subprocess."""
        prompt = self._format_messages(messages)
        response = await self._cli.ask(prompt)
        return ChatResult(
            generations=[ChatGeneration(message=AIMessage(content=response))]
        )

    def with_structured_output(self, schema, *, method=None, include_raw=False, **kwargs):
        """
        Trả về Runnable inject JSON schema vào prompt, parse CLI output thành Pydantic model.

        Thay thế function_calling bằng JSON-in-prompt — hoạt động với mọi method param.
        """
        json_schema = schema.model_json_schema()
        schema_str = json.dumps(json_schema, indent=2, ensure_ascii=False)

        outer = self

        async def _ainvoke(messages, config=None):
            schema_instruction = (
                "\n\n=== OUTPUT FORMAT (BẮT BUỘC) ===\n"
                "Trả lời CHÍNH XÁC bằng 1 JSON object hợp lệ theo schema dưới.\n"
                "KHÔNG markdown, KHÔNG giải thích, KHÔNG code fences. CHỈ raw JSON.\n"
                f"Schema:\n{schema_str}"
            )

            augmented = list(messages)
            if augmented and isinstance(augmented[0], SystemMessage):
                augmented[0] = SystemMessage(
                    content=augmented[0].content + schema_instruction
                )
            else:
                augmented.insert(0, SystemMessage(content=schema_instruction))

            prompt = outer._format_messages(augmented)
            raw = await outer._cli.ask(prompt)
            parsed = GeminiCLI._parse_json(raw)

            try:
                return schema.model_validate(parsed)
            except Exception as e:
                logger.error(
                    "❌ [ChatGeminiCLI] Pydantic validation failed: %s\nRaw: %.500s",
                    e, raw,
                )
                raise

        def _invoke(messages, config=None):
            return asyncio.get_event_loop().run_until_complete(
                _ainvoke(messages, config)
            )

        return RunnableLambda(func=_invoke, afunc=_ainvoke)
