from __future__ import annotations

import json
import os
import re
import sys
import time
import urllib.error
import urllib.request
from dataclasses import dataclass, field
from typing import Any, Iterable

from fileatlas.utils import read_text_file

from . import DISPLAY_NAME, INTERNAL_NAME

try:
    from tqdm import tqdm  # type: ignore
except Exception:  # pragma: no cover
    tqdm = None  # type: ignore[assignment]


# =========================
# Prompt
# =========================

PROMPT = (
    "Role:\n"
    "你是一位拥有深厚代码阅读能力的高级系统架构师。"
    "你擅长通过阅读文件目录、文件摘要和代码片段，快速逆向工程出整个项目的核心架构。\n\n"
    "Context:\n"
    "你将收到两份材料：\n"
    "1. **architecture.md** —— 由 ModuleAtlas 自动生成的模块级架构概览，"
    "包含模块划分、职责摘要、依赖关系图与 fan-in/fan-out 指标。这是你的「全局地图」。\n"
    "2. **project_report.txt（已过滤）** —— 对项目中核心模块下各文件的扫描报告，"
    "包含文件路径、摘要、主要接口和依赖关系。这是你的「文件级证据」。\n\n"
    "请结合这两份材料进行分析。architecture.md 帮你快速定位哪些模块是核心、哪些是辅助；"
    "project_report.txt 提供每个文件的具体细节。\n\n"
    "Task:\n"
    "请分析这些材料，帮我剔除噪音（工具脚本、测试文件、无关配置），"
    "只提取出项目的「脊柱」（Backbone）和「金线」（Golden Thread）。\n"
    "请按照以下结构进行输出：\n\n"
    "第一部分：主线代码文件（The Backbone）\n"
    "请找出支撑项目运行的核心代码链路。不要列出所有文件，只列出决定执行流程的关键节点。"
    "请按逻辑执行顺序排列（例如：从启动入口 -> 调度中心 -> 算法核心 -> 底层Worker），"
    "并为每个文件标注：\n"
    "- 文件路径\n"
    "- 架构角色（例如：Entrypoint, Orchestrator, Algorithm Core, Worker/Engine）\n"
    "- 核心作用（一句话解释它在整个链路中起什么承上启下的作用）\n\n"
    "第二部分：核心贯通示例（The Golden Thread Example）\n"
    "在项目的示例或脚本目录下，找出一个最具代表性的启动脚本。\n"
    "- 选择标准：这个脚本必须能串联起上述的「主线代码」，包含完整的参数配置、模型加载和启动命令。\n"
    "- 解读：请解释为什么选它，并分析它如何激活了上述的主线代码"
    "（例如：它调用了哪个入口？配置了什么核心引擎？）。\n\n"
    "第三部分：逻辑流向图解\n"
    "请用 ASCII 流程图或清晰的层级列表，画出从「核心示例脚本」开始，"
    "数据和控制流是如何流经上述「主线代码」的。\n"
)


@dataclass
class BackboneAgentConfig:
    report_path: str = "project_report.txt"
    out_md_path: str = "backbone_from_report.md"
    # architecture.json / architecture.md produced by ModuleAtlas
    architecture_json_path: str = "architecture.json"
    architecture_md_path: str = "architecture.md"
    # Hard cap LLM input size (chars) to avoid huge prompts.
    max_input_chars: int = 220_000
    # Progress
    show_progress: bool = True
    progress_to_stderr: bool = True


def _p(cfg: BackboneAgentConfig, msg: str) -> None:
    if not bool(getattr(cfg, "show_progress", True)):
        return
    stream = sys.stderr if bool(getattr(cfg, "progress_to_stderr", True)) else sys.stdout
    print(f"[{DISPLAY_NAME}] {msg}", file=stream)


def _wrap_iter(cfg: BackboneAgentConfig, it: Iterable[Any], *, total: int | None, desc: str) -> Iterable[Any]:
    if not bool(getattr(cfg, "show_progress", True)):
        return it
    if tqdm is None:
        return it
    try:
        return tqdm(it, total=total, desc=f"{DISPLAY_NAME}: {desc}", unit="item")  # type: ignore[return-value]
    except Exception:
        return it


def _norm_path(p: str) -> str:
    p = (p or "").replace("\\", "/")
    while p.startswith("./"):
        p = p[2:]
    return p.strip("/")


def _load_llm_from_env_plaintext() -> tuple[str, dict[str, str], str]:
    """
    Returns (provider, settings, model).
    provider:
      - "ollama" (uses /api/chat)
      - "openai_compat" (uses /chat/completions)
    settings includes base_url/api_key when relevant.
    """
    provider = os.environ.get("LLM_PROVIDER", "openai_compat").strip().lower()
    if provider == "ollama":
        return (
            provider,
            {
                "base_url": os.environ.get("OLLAMA_BASE_URL", "http://127.0.0.1:11434").strip(),
            },
            os.environ.get("OLLAMA_MODEL", "qwen3:4b").strip(),
        )
    api_key = os.environ.get("OPENAI_API_KEY", "").strip()
    if not api_key:
        raise RuntimeError("Missing OPENAI_API_KEY (set LLM_PROVIDER=ollama to use Ollama without a key)")
    return (
        "openai_compat",
        {
            "base_url": os.environ.get("OPENAI_BASE_URL", "https://api.openai.com/v1").strip(),
            "api_key": api_key,
        },
        os.environ.get("OPENAI_MODEL", "gpt-4.1-mini").strip(),
    )


def _http_post_json(url: str, payload: dict[str, Any], *, headers: dict[str, str], timeout_s: int) -> dict[str, Any]:
    req = urllib.request.Request(
        url=url,
        method="POST",
        data=json.dumps(payload).encode("utf-8"),
        headers={"Content-Type": "application/json", **(headers or {})},
    )
    try:
        with urllib.request.urlopen(req, timeout=timeout_s) as resp:
            data = resp.read().decode("utf-8", errors="replace")
            obj = json.loads(data)
            return obj if isinstance(obj, dict) else {"raw": obj}
    except urllib.error.HTTPError as e:
        body = ""
        try:
            body = e.read().decode("utf-8", errors="replace")
        except Exception:
            body = ""
        raise RuntimeError(f"LLM HTTPError {e.code}: {body}") from e
    except Exception as e:
        raise RuntimeError(f"LLM request failed: {e}") from e


def call_llm_markdown(*, prompt: str) -> str:
    """
    Call provider with a single user prompt, expecting markdown output.
    (We intentionally do NOT force JSON response_format / format=json.)
    """
    provider, settings, model = _load_llm_from_env_plaintext()
    messages = [{"role": "user", "content": prompt}]

    if provider == "ollama":
        base = settings["base_url"].rstrip("/")
        url = f"{base}/api/chat"
        payload: dict[str, Any] = {
            "model": model,
            "messages": messages,
            "stream": False,
        }
        resp = _http_post_json(url, payload, headers={}, timeout_s=240)
        try:
            return str(resp["message"]["content"])
        except Exception:
            return json.dumps(resp, ensure_ascii=False, indent=2)

    base = settings["base_url"].rstrip("/")
    url = f"{base}/chat/completions"
    payload2: dict[str, Any] = {"model": model, "messages": messages}
    resp2 = _http_post_json(url, payload2, headers={"Authorization": f"Bearer {settings['api_key']}"}, timeout_s=240)
    try:
        return str(resp2["choices"][0]["message"]["content"])
    except Exception:
        return json.dumps(resp2, ensure_ascii=False, indent=2)


_SECTION_RE = re.compile(r"(?m)^##\s+(.+?)\s*$")


def _iter_sections(text: str) -> Iterable[tuple[str, str]]:
    """
    Yield (path, section_text) for each "## path" block.
    """
    if not text:
        return
    matches = list(_SECTION_RE.finditer(text))
    for i, m in enumerate(matches):
        path = _norm_path(m.group(1))
        start = m.start()
        end = matches[i + 1].start() if i + 1 < len(matches) else len(text)
        yield path, text[start:end].rstrip()


# ---------------------------------------------------------------------------
# Architecture-driven noise filtering (replaces hardcoded include/exclude)
# ---------------------------------------------------------------------------

_MODULE_FILTER_PROMPT = (
    "你是软件架构分析师。下面是一个项目中所有模块的摘要信息（由 ModuleAtlas 自动生成）。\n"
    "请判断哪些模块属于项目的「核心执行主线」（Backbone）——即支撑项目核心功能运行的模块。\n\n"
    "排除以下类型的模块：\n"
    "- 纯文档 / CI / 构建配置模块（如 docs、.github、根目录配置文件）\n"
    "- 纯测试模块（如 tests）\n"
    "- 不参与任何代码依赖的孤立模块（fan_in=0 且 fan_out=0）\n\n"
    "保留以下类型的模块：\n"
    "- 项目核心业务逻辑、训练/推理/服务等主线执行模块\n"
    "- 入口/编排/调度模块（即使 fan_in=0，只要 fan_out>0 说明它驱动其他模块）\n"
    "- 被主线模块依赖的关键基础设施模块\n"
    "- 示例/脚本模块（它们可能包含启动入口和「金线」示例）\n\n"
    "注意：fan_in 高的模块可能只是通用工具库，不一定是主线。请根据 purpose 和 summary 的语义来判断，"
    "不要简单按 fan_in/fan_out 数值排序。\n\n"
    "请只返回一个 JSON 数组，包含你认为属于核心主线的 module_id。不要包含任何其他文字。\n"
    "例如: [\"module_a\", \"module_b/sub\"]\n\n"
    "=== 模块列表 ===\n\n"
)


def _load_arch_data(architecture_json_path: str) -> dict[str, Any]:
    """Load and return the parsed architecture.json, or empty dict on failure."""
    try:
        raw = read_text_file(architecture_json_path, max_bytes=50_000_000)
        return json.loads(raw)
    except Exception:
        return {}


def _build_module_summary_table(arch: dict[str, Any]) -> str:
    """
    Build a compact text table of all modules from architecture.json,
    suitable for a lightweight LLM filtering call.
    """
    modules: list[dict[str, Any]] = arch.get("modules") or []
    mg = arch.get("module_graph") or {}
    fan_map: dict[str, dict[str, int]] = {}
    for entry in mg.get("fan_in_out") or []:
        mid = str(entry.get("module_id", ""))
        fan_map[mid] = {
            "fan_in": int(entry.get("fan_in", 0)),
            "fan_out": int(entry.get("fan_out", 0)),
        }

    # Also extract dependency edges for context
    edges: list[dict[str, Any]] = mg.get("edges") or []
    dep_in: dict[str, list[str]] = {}   # module_id -> list of modules that depend on it
    dep_out: dict[str, list[str]] = {}  # module_id -> list of modules it depends on
    for e in edges:
        src = str(e.get("source", ""))
        tgt = str(e.get("target", ""))
        dep_out.setdefault(src, []).append(tgt)
        dep_in.setdefault(tgt, []).append(src)

    lines: list[str] = []
    for mod in modules:
        mid = str(mod.get("module_id", ""))
        purpose = str(mod.get("purpose", ""))
        summary = str(mod.get("summary", ""))
        # Truncate summary to keep the table compact
        if len(summary) > 200:
            summary = summary[:200] + "..."
        fan = fan_map.get(mid, {"fan_in": 0, "fan_out": 0})
        lines.append(f"module_id: {mid}")
        lines.append(f"  purpose: {purpose}")
        lines.append(f"  summary: {summary}")
        lines.append(f"  fan_in: {fan['fan_in']}, fan_out: {fan['fan_out']}")
        dins = dep_in.get(mid, [])
        douts = dep_out.get(mid, [])
        if dins:
            lines.append(f"  depended_by: {dins}")
        if douts:
            lines.append(f"  depends_on: {douts}")
        lines.append("")

    return "\n".join(lines)


def _extract_json_array(text: str) -> list[str]:
    """Extract a JSON array from LLM response text, tolerating markdown fences."""
    # Strip markdown code fences if present
    cleaned = text.strip()
    if cleaned.startswith("```"):
        # Remove opening fence (```json or ```)
        first_newline = cleaned.index("\n") if "\n" in cleaned else len(cleaned)
        cleaned = cleaned[first_newline + 1:]
    if cleaned.endswith("```"):
        cleaned = cleaned[: cleaned.rfind("```")]
    cleaned = cleaned.strip()
    try:
        result = json.loads(cleaned)
        if isinstance(result, list):
            return [str(x) for x in result]
    except json.JSONDecodeError:
        pass
    # Try to find a JSON array anywhere in the text
    match = re.search(r'\[.*?\]', cleaned, re.DOTALL)
    if match:
        try:
            result = json.loads(match.group(0))
            if isinstance(result, list):
                return [str(x) for x in result]
        except json.JSONDecodeError:
            pass
    return []


def _llm_filter_modules(cfg: BackboneAgentConfig, arch: dict[str, Any]) -> set[str] | None:
    """
    Phase 1: Send compact module summaries to LLM, ask it to pick backbone
    candidate modules. Returns a set of module_id strings, or None on failure.
    """
    table = _build_module_summary_table(arch)
    if not table.strip():
        return None

    prompt = _MODULE_FILTER_PROMPT + table
    _p(cfg, f"Phase 1: LLM module filtering (prompt size: {len(prompt)} chars)")

    try:
        response = call_llm_markdown(prompt=prompt)
    except Exception as e:
        _p(cfg, f"Phase 1 LLM call failed: {e}, falling back to fan_in/fan_out rule")
        return None

    selected = _extract_json_array(response)
    if not selected:
        _p(cfg, f"Phase 1: could not parse module list from LLM response, falling back")
        return None

    return {_norm_path(m) for m in selected}


def _fallback_core_modules(arch: dict[str, Any]) -> set[str]:
    """
    Fallback: return modules with fan_in>0 or fan_out>0 (rule-based).
    Used when LLM filtering is unavailable.
    """
    mg = arch.get("module_graph") or {}
    core: set[str] = set()
    for entry in mg.get("fan_in_out") or []:
        mid = _norm_path(str(entry.get("module_id", "")))
        fan_in = int(entry.get("fan_in", 0))
        fan_out = int(entry.get("fan_out", 0))
        if fan_in > 0 or fan_out > 0:
            core.add(mid)
    return core


def _keep_section_by_modules(path: str, core_prefixes: set[str]) -> bool:
    """
    Return True if `path` belongs to any core module prefix.
    If core_prefixes is empty (fallback), keep everything.
    """
    p = _norm_path(path)
    if not p:
        return False
    if not core_prefixes:
        return True
    for prefix in core_prefixes:
        if not prefix:
            continue
        if p == prefix or p.startswith(prefix + "/"):
            return True
    return False


def _select_excerpts(cfg: BackboneAgentConfig, report_text: str, core_prefixes: set[str]) -> str:
    """
    Filter report down to sections belonging to core modules.
    No artificial scoring or re-ordering — let the LLM decide importance.
    Sections are kept in their original order, truncated at max_input_chars.
    """
    out_parts: list[str] = []
    total = 0
    for path, sec in _iter_sections(report_text):
        if not _keep_section_by_modules(path, core_prefixes):
            continue
        chunk = sec + "\n\n"
        if total + len(chunk) > int(cfg.max_input_chars):
            break
        out_parts.append(chunk)
        total += len(chunk)

    return "".join(out_parts).strip()


def generate_backbone_md(cfg: BackboneAgentConfig) -> dict[str, Any]:
    # --- Load architecture data ---
    _p(cfg, f"reading architecture: {cfg.architecture_json_path}")
    arch = _load_arch_data(cfg.architecture_json_path)

    # --- Phase 1: LLM-based module filtering ---
    filter_method: str
    if arch:
        core_prefixes = _llm_filter_modules(cfg, arch)
        if core_prefixes is not None:
            filter_method = "llm"
            _p(cfg, f"Phase 1 result — LLM selected {len(core_prefixes)} modules: {sorted(core_prefixes)}")
        else:
            core_prefixes = _fallback_core_modules(arch)
            filter_method = "fan_in_out_rule"
            _p(cfg, f"Fallback to fan_in/fan_out rule — {len(core_prefixes)} modules: {sorted(core_prefixes)}")
    else:
        core_prefixes = set()
        filter_method = "none"
        _p(cfg, "WARNING: no architecture.json found, keeping all report sections")

    # --- Load architecture.md for context ---
    arch_md = ""
    try:
        arch_md = read_text_file(cfg.architecture_md_path, max_bytes=5_000_000)
        _p(cfg, f"loaded architecture.md: {len(arch_md)} chars")
    except Exception:
        _p(cfg, f"WARNING: could not read {cfg.architecture_md_path}, proceeding without it")

    # --- Phase 2: Filter report and call main LLM ---
    _p(cfg, f"reading report: {cfg.report_path}")
    txt = read_text_file(cfg.report_path, max_bytes=25_000_000)
    _p(cfg, "filtering report (using module selection from Phase 1)")
    excerpts = _select_excerpts(cfg, txt, core_prefixes)
    _p(cfg, f"excerpt size: {len(excerpts)} chars (cap={cfg.max_input_chars})")

    # --- Assemble prompt ---
    parts = [PROMPT]
    if arch_md:
        parts.append("\n\n=== architecture.md ===\n\n")
        parts.append(arch_md)
    parts.append("\n\n=== project_report.txt (filtered excerpts) ===\n\n")
    parts.append(excerpts)
    parts.append("\n")
    full_prompt = "".join(parts)

    _p(cfg, f"Phase 2: total prompt size: {len(full_prompt)} chars")
    _p(cfg, "Phase 2: calling LLM (markdown output)")
    md: str
    llm_error: str | None = None
    try:
        md = call_llm_markdown(prompt=full_prompt)
    except Exception as e:
        llm_error = str(e)
        md = (
            "## BackboneAgent\n\n"
            "LLM call failed, so no analysis was generated.\n\n"
            f"- **Error**: `{llm_error}`\n\n"
            "### How to fix\n\n"
            "- **Ollama (no key required)**:\n\n"
            "```bash\n"
            "export LLM_PROVIDER=ollama\n"
            "export OLLAMA_BASE_URL=http://127.0.0.1:11434\n"
            "export OLLAMA_MODEL=qwen3:4b\n"
            "python -m backboneagent --report project_report.txt --out backbone_from_report.md\n"
            "```\n\n"
            "- **OpenAI-compatible**:\n\n"
            "```bash\n"
            "export OPENAI_API_KEY=...\n"
            "export OPENAI_MODEL=...\n"
            "export OPENAI_BASE_URL=...\n"
            "python -m backboneagent --report project_report.txt --out backbone_from_report.md\n"
            "```\n"
        )

    # Write output
    out_path = str(cfg.out_md_path)
    with open(out_path, "w", encoding="utf-8") as f:
        f.write(md if md.endswith("\n") else (md + "\n"))

    meta = {
        "generated_at": time.strftime("%Y-%m-%d %H:%M:%S"),
        "agent": INTERNAL_NAME,
        "agent_display_name": DISPLAY_NAME,
        "report_path": os.path.abspath(cfg.report_path),
        "out_md_path": os.path.abspath(out_path),
        "architecture_json_path": os.path.abspath(cfg.architecture_json_path),
        "architecture_md_path": os.path.abspath(cfg.architecture_md_path),
        "filter_method": filter_method,
        "core_modules": sorted(core_prefixes),
        "excerpt_chars": int(len(excerpts)),
        "max_input_chars": int(cfg.max_input_chars),
        "llm_error": llm_error or "",
    }
    _p(cfg, "done")
    return {"meta": meta, "out_md_path": out_path, "markdown": md}
