"""
Agentic code generator for SkyDiscover.

Gives the LLM a multi-turn tool-calling loop so it can read files and search
the codebase before outputting the final answer.  Plugs into DiscoveryManager
as a drop-in replacement for the single-shot ``llm.generate()`` call.

Tools: read_file, search.  Extract the final answer from the response.
"""

import asyncio
import concurrent.futures
import fnmatch
import json
import logging
import os
import re
import time
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple
from skydiscover.llm.openai import is_openai_reasoning_model, GOOGLE_AI_STUDIO_DOMAIN
from skydiscover.utils.repo_map import build_repo_map

logger = logging.getLogger(__name__)


# Tool schemas (loaded from external JSON file)
_TOOL_SCHEMAS_PATH = Path(__file__).parent / "tool_schemas" / "agentic_tools.json"

with open(_TOOL_SCHEMAS_PATH, "r") as _f:
    TOOL_SCHEMAS = json.load(_f)

# Prompt template
_AGENTIC_PROMPT_PATH = Path(__file__).parent.parent / "prompt" / "default" / "templates" / "agentic_system_message.txt"

with open(_AGENTIC_PROMPT_PATH, "r") as _f:
    _AGENTIC_SYSTEM_PROMPT = _f.read()


class AgenticGenerator:
    """
    Simple v0 tool-calling agent that explores a codebase before generating code.

    Replaces the single-shot LLM call with a multi-turn loop.  The agent
    can ``read_file`` and ``search``.  When it stops calling tools and
    responds with text, that text is treated as its final output.
    """

    def __init__(self, llm_pool, config):
        self.llm_pool = llm_pool
        self.config = config

    async def generate(
        self,
        system_message: str,
        user_message: str,
    ) -> Optional[str]:
        """
        Run the agent loop.

        Returns the submitted code string, or *None* if the agent did not
        produce output (caller should fall back to direct generation).
        """

        cfg = self.config
        files_read: set = set()
        conversation: List[Dict[str, Any]] = []
        t0 = time.time()

        # Build the initial messages
        sys_prompt = self._system_prompt(system_message)
        repo_map = self._repo_map()
        user_parts = [user_message]
        if repo_map:
            user_parts.append(f"\n## Project structure\n```\n{repo_map}\n```")
        conversation.append({"role": "user", "content": "\n".join(user_parts)})

        for step in range(cfg.max_steps):
            # Guard: overall timeout
            if time.time() - t0 > cfg.overall_timeout:
                logger.warning("Agentic loop timed out at step %d", step)
                break

            # Guard: context size
            ctx_chars = _context_chars(sys_prompt, conversation)
            if ctx_chars > cfg.max_context_chars:
                logger.warning("Context limit hit (%d chars), requesting code", ctx_chars)
                conversation.append({
                    "role": "user",
                    "content": (
                        "Context limit reached. Please output your complete "
                        "improved program now."
                    ),
                })

            # Call LLM
            try:
                assistant_msg = await asyncio.wait_for(
                    self._call_llm(sys_prompt, conversation),
                    timeout=cfg.per_step_timeout,
                )
            except asyncio.TimeoutError:
                logger.warning("Step %d: LLM call timed out", step)
                conversation.append({
                    "role": "user",
                    "content": "Response timed out. Please output your solution or try a simpler action.",
                })
                continue
            except Exception as e:
                logger.error("Step %d: LLM error: %s", step, e)
                break

            tool_calls = assistant_msg.get("tool_calls", [])
            text_content = assistant_msg.get("content", "").strip()
            conversation.append(assistant_msg)

            if not tool_calls:
                if text_content:
                    logger.info(
                        "Agent responded with text at step %d (%d files read)",
                        step, len(files_read),
                    )
                    return text_content
                # Empty response with no tools — nudge
                conversation.append({
                    "role": "user",
                    "content": (
                        "Please either use a tool to explore the codebase, "
                        "or output your complete improved program."
                    ),
                })
                continue

            for tc in tool_calls:
                fn = tc.get("function", {})
                name = fn.get("name", "")
                raw = fn.get("arguments", "{}")
                tc_id = tc.get("id", "")

                # Parse args
                try:
                    args = json.loads(raw)
                except (json.JSONDecodeError, TypeError) as e:
                    conversation.append({
                        "role": "tool",
                        "tool_call_id": tc_id,
                        "content": f"Bad JSON: {e}. Raw: {raw[:200]}",
                    })
                    continue

                # Dispatch
                logger.info(
                    "Step %d: tool=%s args=%s", step, name,
                    {k: v[:80] if isinstance(v, str) and len(v) > 80 else v
                     for k, v in args.items()},
                )
                result = self._run_tool(name, args, files_read)

                conversation.append({
                    "role": "tool",
                    "tool_call_id": tc_id,
                    "content": result["content"],
                })

        logger.warning("Agent loop ended without producing code")
        return None

    def _system_prompt(self, original: str) -> str:
        """Append the agentic instructions to the base system message."""
        return f"{original}\n\n{_AGENTIC_SYSTEM_PROMPT}"

    def _repo_map(self) -> str:
        """Depth-limited directory tree of the codebase.

        Shows the agent what files are available so it can choose what to
        ``read_file`` without blind guessing.  See :func:`build_repo_map`.
        """
        return build_repo_map(
            self.config.codebase_root,
            max_depth=self.config.repo_map_max_depth,
            allowed_extensions=self.config.allowed_extensions,
            excluded_dirs=self.config.excluded_dirs,
        )

    async def _call_llm(
        self, system_message: str, conversation: List[Dict[str, Any]]
    ) -> Dict[str, Any]:
        """Call the LLM with tool schemas via the OpenAI client directly."""

        # Sample model (same logic as LLMPool.generate)
        idx = self.llm_pool.random_state.choices(
            range(len(self.llm_pool.models)),
            weights=self.llm_pool.weights,
            k=1,
        )[0]
        model = self.llm_pool.models[idx]
        model_name = getattr(model, "model", type(model).__name__)
        logger.info("Agentic call: model %s", model_name)

        if not hasattr(model, "client"):
            raise RuntimeError(
                f"Agentic mode needs an OpenAI-compatible LLM. "
                f"{type(model).__name__} has no .client attribute."
            )

        messages = [{"role": "system", "content": system_message}] + conversation

        api_base = getattr(model, "api_base", "") or ""
        is_reasoning = is_openai_reasoning_model(model.model, api_base)

        params = {
            "model": model.model,
            "messages": messages,
            "tools": TOOL_SCHEMAS,
            "tool_choice": "auto",
        }
        if is_reasoning:
            if getattr(model, "max_tokens", None) is not None:
                params["max_completion_tokens"] = model.max_tokens
            reasoning_effort = getattr(model, "reasoning_effort", None)
            if reasoning_effort is not None:
                params["reasoning_effort"] = reasoning_effort
        else:
            if getattr(model, "temperature", None) is not None:
                params["temperature"] = model.temperature
            if getattr(model, "top_p", None) is not None:
                params["top_p"] = model.top_p
            if getattr(model, "max_tokens", None) is not None:
                params["max_tokens"] = model.max_tokens

        seed = getattr(model, "random_seed", None)
        if seed is not None:
            if GOOGLE_AI_STUDIO_DOMAIN not in api_base.lower():
                params["seed"] = seed

        loop = asyncio.get_running_loop()
        resp = await loop.run_in_executor(
            None, lambda: model.client.chat.completions.create(**params)
        )

        msg = resp.choices[0].message
        out: Dict[str, Any] = {"role": "assistant", "content": msg.content or ""}
        if msg.tool_calls:
            out["tool_calls"] = [
                {
                    "id": tc.id,
                    "type": "function",
                    "function": {
                        "name": tc.function.name,
                        "arguments": tc.function.arguments,
                    },
                }
                for tc in msg.tool_calls
            ]
        return out

    # ------------------------------------------------------------------
    # Tool dispatch & implementations
    # ------------------------------------------------------------------

    def _run_tool(
        self, name: str, args: Dict[str, Any], files_read: set
    ) -> Dict[str, Any]:
        """Execute a single tool call. All tools are synchronous."""
        try:
            if name == "read_file":
                return self._tool_read_file(args, files_read)
            elif name == "search":
                return self._tool_search(args)
            else:
                return _err(
                    f"Unknown tool '{name}'. Available: read_file, search."
                )
        except Exception as e:
            return _err(f"Tool '{name}' raised: {e}")

    def _tool_read_file(self, args: Dict[str, Any], files_read: set) -> Dict[str, Any]:
        path = args.get("path", "")
        if not path:
            return _err("'path' is required.")

        root = self.config.codebase_root
        if not root:
            return _err("'agentic.codebase_root' is not configured.")
        full = os.path.join(root, path) if not os.path.isabs(path) else path

        ok, resolved, err = _validate_path(
            full, root, self.config.allowed_extensions, self.config.excluded_dirs
        )
        if not ok:
            return _err(err)

        # Read limit
        if resolved not in files_read and len(files_read) >= self.config.max_files_read:
            return _err(
                f"Read limit ({self.config.max_files_read}). "
                f"You have enough context — please output your solution."
            )

        try:
            with open(resolved, "r", encoding="utf-8", errors="replace") as f:
                lines = f.readlines()
        except Exception as e:
            return _err(f"Cannot read: {e}")

        total = len(lines)
        raw_start = args.get("line_start")
        raw_end = args.get("line_end")
        start = max(1, int(raw_start) if raw_start is not None else 1) - 1
        end = min(total, int(raw_end) if raw_end is not None else total)
        content = "".join(lines[start:end])

        # Truncate large files
        if len(content) > self.config.max_file_chars:
            half = self.config.max_file_chars // 2
            content = (
                content[:half]
                + f"\n\n... ({len(content) - self.config.max_file_chars} chars truncated, "
                f"{total} lines total) ...\n\n"
                + content[-half:]
            )

        files_read.add(resolved)
        rel = os.path.relpath(resolved, root)

        numbered = []
        for i, ln in enumerate(content.splitlines(True), start=start + 1):
            line = ln.rstrip("\n")
            numbered.append(f"{i:4d} | {line}")

        return {"content": f"{rel} (lines {start + 1}-{end} of {total})\n" + "\n".join(numbered)}

    def _tool_search(self, args: Dict[str, Any]) -> Dict[str, Any]:
        pattern = args.get("pattern", "")
        glob_pat = args.get("file_glob", "*.py")

        if not pattern:
            return _err("'pattern' is required.")
        if len(pattern) > self.config.max_regex_length:
            return _err(f"Pattern too long ({len(pattern)} > {self.config.max_regex_length}).")

        safety_err = _check_regex_safety(pattern)
        if safety_err:
            return _err(safety_err)

        try:
            compiled = re.compile(pattern)
        except re.error as e:
            return _err(f"Invalid regex: {e}")

        root = self.config.codebase_root
        if not root:
            return _err("'agentic.codebase_root' is not configured.")
        excluded = set(self.config.excluded_dirs)
        allowed = set(self.config.allowed_extensions)
        matches: List[str] = []
        n_files = 0

        for dirpath, dirnames, filenames in os.walk(root):
            dirnames[:] = [d for d in dirnames if not d.startswith(".") and d not in excluded]

            for fname in filenames:
                if not fnmatch.fnmatch(fname, glob_pat):
                    continue
                if os.path.splitext(fname)[1].lower() not in allowed:
                    continue

                fpath = os.path.join(dirpath, fname)
                try:
                    size = os.path.getsize(fpath)
                    if size > self.config.max_file_chars:
                        continue  # skip very large files
                    with open(fpath, "r", encoding="utf-8", errors="replace") as f:
                        text = f.read()
                except Exception:
                    continue

                n_files += 1
                ok, file_hits, err = _safe_regex_search(compiled, text, self.config.regex_timeout)
                if not ok:
                    return _err(err)

                rel = os.path.relpath(fpath, root)
                for hit in file_hits:
                    matches.append(f"{rel}:{hit}")
                    if len(matches) >= self.config.max_search_results:
                        break
                if len(matches) >= self.config.max_search_results:
                    break
            if len(matches) >= self.config.max_search_results:
                break

        if not matches:
            return {"content": f"No matches for '{pattern}' in {n_files} files."}

        cap = ""
        if len(matches) >= self.config.max_search_results:
            cap = f"\n(capped at {self.config.max_search_results} results)"

        return {"content": "\n".join(matches) + cap}


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _err(msg: str) -> Dict[str, Any]:
    return {"content": msg, "_error": True}


def _context_chars(system: str, conversation: List[Dict[str, Any]]) -> int:
    n = len(system)
    for msg in conversation:
        n += len(msg.get("content", ""))
        for tc in msg.get("tool_calls", []):
            n += len(tc.get("function", {}).get("arguments", ""))
    return n


_SENSITIVE_FILENAMES = frozenset({
    ".env", ".env.local", ".env.production", ".env.staging",
    "secrets.json", "secrets.yaml", "secrets.yml",
    "credentials.json", "credentials.yaml",
    "service-account.json", "service_account.json",
    ".netrc", ".pgpass", ".my.cnf",
})


def _validate_path(
    requested: str, root: str, allowed_ext: tuple, excluded_dirs: tuple
) -> Tuple[bool, str, str]:
    """Resolve and validate a file path. Returns (ok, resolved, error)."""
    try:
        resolved = os.path.realpath(requested)
    except (OSError, ValueError) as e:
        return False, "", f"Invalid path: {e}"

    root_abs = os.path.realpath(root)
    if not resolved.startswith(root_abs + os.sep) and resolved != root_abs:
        return False, "", "Path outside codebase root."

    # Excluded dirs
    try:
        rel = os.path.relpath(resolved, root_abs)
        for part in Path(rel).parts:
            if part in excluded_dirs:
                return False, "", f"Path in excluded directory '{part}'."
    except ValueError:
        # On Windows, relpath raises ValueError for paths on different drives.
        # The startswith check above already guards against escape, so this is safe to skip.
        pass

    # Block sensitive filenames regardless of whether they exist
    basename = os.path.basename(resolved).lower()
    if basename in _SENSITIVE_FILENAMES:
        return False, "", f"Access denied: '{basename}' may contain secrets."

    if not os.path.isfile(resolved):
        # Show siblings so the agent can self-correct
        parent_dir = os.path.dirname(resolved)
        if os.path.isdir(parent_dir):
            try:
                siblings = sorted(os.listdir(parent_dir))[:15]
                rel_dir = os.path.relpath(parent_dir, root_abs)
                return False, "", (
                    f"Not found: '{os.path.basename(resolved)}'. "
                    f"'{rel_dir}/' contains: {siblings}"
                )
            except OSError:
                pass
        return False, "", f"File not found: '{requested}'."

    ext = os.path.splitext(resolved)[1].lower()
    if ext not in allowed_ext:
        return False, "", f"Extension '{ext}' not allowed."

    return True, resolved, ""


# Nested quantifiers cause catastrophic backtracking: (a+)+, (.*)+, (\w+)*, etc.
_NESTED_QUANTIFIER_RE = re.compile(
    r"\([^)]*[+*][^)]*\)\s*[+*?]"
    r"|"
    r"\([^)]*[+*][^)]*\)\s*\{",
)


def _check_regex_safety(pattern: str) -> Optional[str]:
    """Reject patterns with constructs that cause catastrophic backtracking."""
    if _NESTED_QUANTIFIER_RE.search(pattern):
        return (
            "Pattern has nested quantifiers (e.g. '(a+)+') which cause catastrophic "
            "backtracking. Use a simpler pattern."
        )
    return None


# Max chars per line to search — very long lines (e.g. minified code)
# amplify backtracking even for simple patterns.
_MAX_SEARCH_LINE_LEN = 2000


def _safe_regex_search(
    compiled: "re.Pattern", text: str, timeout: float = 2.0
) -> Tuple[bool, List[str], str]:
    """
    Safe regex search with a thread timeout as best effort.
    """
    def do_search():
        hits = []
        for i, line in enumerate(text.splitlines(), 1):
            if len(line) > _MAX_SEARCH_LINE_LEN:
                continue
            if compiled.search(line):
                hits.append(f"{i}: {line}")
        return hits

    pool = concurrent.futures.ThreadPoolExecutor(max_workers=1)
    fut = pool.submit(do_search)
    try:
        result = fut.result(timeout=timeout)
        pool.shutdown(wait=False)
        return True, result, ""
    except concurrent.futures.TimeoutError:
        pool.shutdown(wait=False)
        return False, [], f"Regex timed out ({timeout}s). Simplify the pattern."
