from __future__ import annotations

from concurrent.futures import ThreadPoolExecutor, as_completed
import json
import os
import time
from dataclasses import dataclass
from typing import Any

from tqdm import tqdm

from .llm import extract_json_content, extract_json_content_ollama, load_llm_from_env
from . import DISPLAY_NAME, INTERNAL_NAME
from .schema import FileAnalysis, FileFacts, FileRecord, ProjectMap
from .tools import ASTParser, ChunkSlicer, FileWalker, RegexSkeletonParser, SimpleIndexer
from .utils import estimate_tokens, read_text_file, sha256_file


@dataclass
class FileAtlasConfig:
    root: str = "."
    out_path: str = "project_map.json"
    use_gitignore: bool = True
    include_non_python_skeleton: bool = True
    build_index: bool = False
    index_out_path: str = "project_index.json"
    # optional human-readable report
    txt_out_path: str | None = None
    # chunking
    max_tokens_per_chunk: int = 8000
    # caching
    use_cache: bool = True
    # LLM analysis (always on; best-effort). Tune concurrency to match your vLLM capacity.
    llm_concurrency: int = 4
    llm_max_excerpt_tokens: int = 1500
    # if True, re-run LLM analysis even when cached analysis exists
    llm_refresh: bool = False


def _load_existing_project_map(path: str) -> ProjectMap | None:
    try:
        with open(path, "r", encoding="utf-8") as f:
            return json.load(f)
    except Exception:
        return None


def _default_analysis() -> FileAnalysis:
    return {
        "summary": "",
        "responsibilities": [],
        "external_dependencies": [],
        "key_interfaces": [],
        "risks": [],
    }


def _llm_analyze_file(
    *,
    path: str,
    language: str,
    facts: FileFacts,
    code: str | None,
    max_excerpt_tokens: int,
) -> FileAnalysis:
    """
    LLM phase: only for "opinions" (responsibility, key interface semantics, risks).
    Deterministic facts come from parsers and must not be invented here.
    """
    provider, client, model = load_llm_from_env()

    import_modules = []
    for imp in facts.get("imports", []) or []:
        mod = (imp.get("module") or "").strip()
        if mod:
            import_modules.append(mod)
    import_modules = sorted(set(import_modules))

    sym_lines: list[str] = []
    for s in facts.get("symbols", []) or []:
        sig = s.get("signature") or ""
        qn = s.get("qualname") or s.get("name") or ""
        sp = s.get("span") or {}
        if qn and sig:
            sym_lines.append(f"- {qn} @ L{sp.get('start_line','?')}-L{sp.get('end_line','?')}: {sig}")
        elif qn:
            sym_lines.append(f"- {qn}")

    excerpt = ""
    if code:
        est = estimate_tokens(code)
        if est.approx_tokens <= max_excerpt_tokens:
            excerpt = code
        else:
            # Use head+tail so entry-guards at file bottom aren't always missed.
            approx_chars = max_excerpt_tokens * 4
            half = max(200, approx_chars // 2)
            head = code[:half]
            tail = code[-half:] if len(code) > half else ""
            excerpt = head + "\n\n# ... truncated ...\n\n" + tail

    user_payload = {
        "path": path,
        "language": language,
        "imports": import_modules,
        "symbols": sym_lines,
        "excerpt": excerpt,
    }

    system = (
        "你是资深软件架构师与代码分析专家。你的核心能力是透过代码细节洞察单个代码文件在整个系统中的真实角色与职责。\n"
        "你的任务是对输入的【单个代码文件】进行深度语义分析，并输出一份极简但高密度的信息报告。\n"
        "你必须严格遵循用户提供的分析要求与 JSON Schema，不得添加任何额外解释或自然语言说明。\n\n"
    )
    prompt = (
        "请分析以下代码文件的上下文，并严格按照下方的 JSON Schema 输出分析结果。\n\n"
        "### 分析要求 (Analysis Guidelines)\n\n"
        "1. Summary:\n"
        "- 必须精炼且一针见血，在3-4句话左右\n"
        "- 不要使用“这是一个包含……”，“该文件定义了...”这类描述\n"
        "- 直接说明该文件在系统中的真实作用\n"
        "2. Responsibilities:\n"
        "- 提取 3–7 个核心职责\n"
        "- 关注业务 / 系统语义\n"
        "- 避免描述语法或实现细节（如 if / for / try）\n\n"
        "3. Key Interfaces:\n"
        "- 列出该文件对外暴露或被依赖的关键接口\n"
        "- 优先使用符号名 / qualname\n"
        "- 例如：Trainer.train / Agent.act / encrypt_password\n\n"
        "### Output Format (JSON Only)\n"
        "必须严格按以下 Schema 输出，不得包含任何额外文本：\n\n"
        "{\n"
        '  "summary": "string",\n'
        '  "responsibilities": ["string", "string", "..."],\n'
        '  "external_dependencies": ["string", "..."],\n'
        '  "key_interfaces": ["string", "..."],\n'
        "}\n\n"
        "external_dependencies: 0-12 条字符串（只列 INPUT 里出现过/可推断的外部依赖）\n\n"
        "INPUT CODE CONTEXT:\n"
    )

    messages = [
        {"role": "system", "content": system},
        {"role": "user", "content": prompt + "\n\nINPUT:\n" + json.dumps(user_payload, ensure_ascii=False, indent=2)},
    ]

    if provider == "ollama":
        resp = client.chat(model=model, messages=messages)  # type: ignore[attr-defined]
        obj = extract_json_content_ollama(resp)
    else:
        resp = client.chat_completions(model=model, messages=messages)  # type: ignore[attr-defined]
        obj = extract_json_content(resp)
    if not isinstance(obj, dict):
        return {"error": "bad_llm_json", "raw": obj}
    # Propagate structured extraction failures (e.g., bad_response/non_json_content).
    if "error" in obj and not any(k in obj for k in ["summary", "responsibilities", "external_dependencies", "key_interfaces", "risks"]):
        return {"error": str(obj.get("error")), "raw": obj.get("raw", obj)}

    out = _default_analysis()
    for k in ["summary", "responsibilities", "external_dependencies", "key_interfaces", "risks"]:
        if k in obj:
            out[k] = obj[k]  # type: ignore[index]
    return out


def _analysis_is_retryable(analysis: FileAnalysis) -> bool:
    """
    Conservative retry policy: only retry when we clearly failed to get a usable JSON object,
    or when the request itself failed.
    """
    err = analysis.get("error") or ""
    if not isinstance(err, str) or not err:
        return False
    # Non-retryable: local configuration issues (will never succeed by retrying).
    if "Missing OPENAI_API_KEY" in err:
        return False
    if err in {"bad_llm_json", "bad_response", "non_json_content"}:
        return True
    if err.startswith("LLMError:") or err.startswith("LLM HTTPError"):
        return True
    if err.startswith("Ollama HTTPError") or err.startswith("Ollama request failed"):
        return True
    if err.startswith("LLM request failed"):
        return True
    return False


def _llm_analyze_file_with_retry(
    *,
    path: str,
    language: str,
    facts: FileFacts,
    code: str | None,
    max_excerpt_tokens: int,
    max_attempts: int = 3,
    backoff_seconds: float = 0.5,
) -> FileAnalysis:
    """
    Best-effort wrapper so transient vLLM/OpenAI-compat failures don't kill the whole run.
    """
    attempts = max(1, int(max_attempts))
    backoff = max(0.0, float(backoff_seconds))
    last: FileAnalysis = {"error": "LLMError: unknown", "raw": None}
    for i in range(attempts):
        try:
            analysis = _llm_analyze_file(
                path=path,
                language=language,
                facts=facts,
                code=code,
                max_excerpt_tokens=max_excerpt_tokens,
            )
        except Exception as e:
            analysis = {"error": f"LLMError: {e}", "raw": None}
        last = analysis
        if i >= attempts - 1 or not _analysis_is_retryable(analysis):
            return analysis
        time.sleep(backoff * (2**i))
    return last


def build_project_map(cfg: FileAtlasConfig) -> ProjectMap:
    root = os.path.abspath(cfg.root)

    walker = FileWalker(root=root, use_gitignore=cfg.use_gitignore)
    parser = ASTParser()
    regex_parser = RegexSkeletonParser()
    slicer = ChunkSlicer(max_tokens_per_chunk=cfg.max_tokens_per_chunk)
    indexer = SimpleIndexer()

    t0 = time.time()
    existing = _load_existing_project_map(cfg.out_path) if cfg.use_cache else None
    existing_files = (existing or {}).get("files", {}) if isinstance(existing, dict) else {}

    # Phase: scan paths (fast)
    candidates: list[FileRecord] = walker.walk()

    # Phase: hash with progress (enables cache + gives user feedback)
    records: list[FileRecord] = []
    for rec in tqdm(candidates, desc=f"{DISPLAY_NAME}: hashing", unit="file"):
        abspath = os.path.join(root, rec.path.replace("/", os.sep))
        try:
            sha = sha256_file(abspath)
        except Exception:
            sha = ""
        records.append(
            FileRecord(
                path=rec.path,
                language=rec.language,
                size_bytes=rec.size_bytes,
                sha256=sha,
            )
        )

    project: ProjectMap = {
        "meta": {
            "root": root,
            "generated_at": time.strftime("%Y-%m-%d %H:%M:%S"),
            "file_count": len(records),
            "agent": INTERNAL_NAME,
            "agent_display_name": DISPLAY_NAME,
            "version": "0.1.1",
            "use_cache": bool(cfg.use_cache),
            "llm_enabled": True,
            "llm_concurrency": max(1, int(cfg.llm_concurrency)),
            "llm_refresh": bool(cfg.llm_refresh),
            "llm_max_excerpt_tokens": int(cfg.llm_max_excerpt_tokens),
        },
        "files": {},
        "symbols": {},
        "edges": {"calls": [], "imports": []},
    }

    # Phase: parse files
    for rec in tqdm(records, desc=f"{DISPLAY_NAME}: parsing", unit="file"):
        # cache hit?
        if cfg.use_cache:
            prev = existing_files.get(rec.path) if isinstance(existing_files, dict) else None
            try:
                prev_sha = prev.get("facts", {}).get("sha256")  # type: ignore[union-attr]
            except Exception:
                prev_sha = None
            if prev and prev_sha and prev_sha == rec.sha256:
                project["files"][rec.path] = prev
                facts = prev.get("facts", {}) if isinstance(prev, dict) else {}
                for sym in facts.get("symbols", []) or []:
                    sid = sym.get("symbol_id")
                    if sid:
                        project["symbols"][sid] = sym
                for call in facts.get("calls", []) or []:
                    project["edges"]["calls"].append(call)
                for imp in facts.get("imports", []) or []:
                    project["edges"]["imports"].append({"path": rec.path, **imp})
                continue

        abspath = os.path.join(root, rec.path.replace("/", os.sep))
        try:
            code = read_text_file(abspath)
        except OSError as e:
            project["files"][rec.path] = {
                "facts": {
                    "path": rec.path,
                    "language": rec.language,
                    "size_bytes": rec.size_bytes,
                    "sha256": rec.sha256,
                    "parse_errors": [f"ReadError: {e}"],
                }
            }
            continue

        facts: FileFacts
        if rec.language == "python":
            facts = parser.parse_python(code, rec.path)
        elif cfg.include_non_python_skeleton and rec.language in {"javascript", "typescript"}:
            facts = regex_parser.parse_js_ts(code, rec.path, rec.language)
        else:
            facts = {
                "path": rec.path,
                "language": rec.language,
                "imports": [],
                "symbols": [],
                "calls": [],
                "parse_errors": [],
            }

        # attach file metadata
        facts["size_bytes"] = rec.size_bytes
        facts["sha256"] = rec.sha256

        # chunk slicing (always useful for later evidence retrieval)
        try:
            chunks = slicer.slice_by_symbols(path=rec.path, code=code, symbols=facts.get("symbols", []))  # type: ignore[arg-type]
            # Store only chunk spans/ids; not storing raw text to keep project_map small by default.
            chunk_meta = [
                {
                    "symbol_id": c.symbol_id,
                    "path": c.path,
                    "start_line": c.start_line,
                    "end_line": c.end_line,
                    "approx_tokens": c.approx_tokens,
                }
                for c in chunks
            ]
        except Exception as e:
            chunk_meta = [{"error": f"ChunkError: {e}"}]

        project["files"][rec.path] = {"facts": facts, "chunks": chunk_meta}

        # global symbols map
        for sym in facts.get("symbols", []) or []:
            sid = sym.get("symbol_id")
            if not sid:
                continue
            project["symbols"][sid] = sym  # type: ignore[index]

        # edges: calls
        for call in facts.get("calls", []) or []:
            project["edges"]["calls"].append(call)

        # edges: imports
        for imp in facts.get("imports", []) or []:
            project["edges"]["imports"].append({"path": rec.path, **imp})

    # Phase: LLM analysis (always on; parallelized). Deterministic facts come from parsers.
    to_analyze: list[tuple[FileRecord, FileFacts]] = []
    for rec in records:
        file_entry = project["files"].get(rec.path, {})
        if not isinstance(file_entry, dict):
            continue
        # If cached entry already has analysis, keep it unless user forces refresh.
        if "analysis" in file_entry and not cfg.llm_refresh:
            continue
        facts = file_entry.get("facts", {})
        if not isinstance(facts, dict):
            continue
        to_analyze.append((rec, facts))  # type: ignore[arg-type]

    if to_analyze:
        max_workers = max(1, int(cfg.llm_concurrency))

        def _worker(rec: FileRecord, facts: FileFacts) -> tuple[str, FileAnalysis]:
            abspath = os.path.join(root, rec.path.replace("/", os.sep))
            code: str | None = None
            try:
                code = read_text_file(abspath)
            except Exception:
                code = None
            analysis = _llm_analyze_file_with_retry(
                path=rec.path,
                language=rec.language,
                facts=facts,
                code=code,
                max_excerpt_tokens=cfg.llm_max_excerpt_tokens,
            )
            return rec.path, analysis

        with tqdm(
            total=len(to_analyze),
            desc=f"{DISPLAY_NAME}: LLM analysis ({max_workers} workers)",
            unit="file",
        ) as pbar:
            with ThreadPoolExecutor(max_workers=max_workers) as ex:
                futures = [ex.submit(_worker, rec, facts) for rec, facts in to_analyze]
                for fut in as_completed(futures):
                    try:
                        path, analysis = fut.result()
                    except Exception as e:
                        # Should be rare; keep moving and record nothing for this file.
                        _ = e
                        pbar.update(1)
                        continue
                    file_entry = project["files"].get(path)
                    if isinstance(file_entry, dict):
                        file_entry["analysis"] = analysis
                    pbar.update(1)

    # Optional index
    if cfg.build_index:
        inv = indexer.build_inverted_index(symbols=project["symbols"], files=project["files"])
        project["meta"]["index"] = {"type": "inverted_json", "path": cfg.index_out_path}
        indexer.save(inv, os.path.join(root, cfg.index_out_path))

    project["meta"]["elapsed_seconds"] = round(time.time() - t0, 3)
    return project


def save_project_map(project: ProjectMap, out_path: str) -> None:
    os.makedirs(os.path.dirname(os.path.abspath(out_path)) or ".", exist_ok=True)
    with open(out_path, "w", encoding="utf-8") as f:
        json.dump(project, f, ensure_ascii=False, indent=2, sort_keys=True)


def save_txt_report(project: ProjectMap, out_path: str) -> None:
    """
    Optional: a compact, human-readable report for quick review.
    """
    lines: list[str] = []
    meta = project.get("meta", {}) if isinstance(project, dict) else {}
    lines.append(f"{DISPLAY_NAME} Report")
    lines.append(f"root: {meta.get('root','')}")
    lines.append(f"generated_at: {meta.get('generated_at','')}")
    lines.append(f"file_count: {meta.get('file_count','')}")
    lines.append("")

    files = project.get("files", {}) if isinstance(project, dict) else {}
    if isinstance(files, dict):
        for path in sorted(files.keys()):
            entry = files.get(path) or {}
            facts = entry.get("facts", {}) if isinstance(entry, dict) else {}
            analysis = entry.get("analysis", {}) if isinstance(entry, dict) else {}
            lines.append(f"## {path}")
            if isinstance(facts, dict):
                lines.append(f"- language: {facts.get('language','')}")
                lines.append(f"- sha256: {facts.get('sha256','')}")
                syms = facts.get('symbols', []) or []
                lines.append(f"- symbols: {len(syms) if isinstance(syms, list) else 0}")
                imps = facts.get('imports', []) or []
                lines.append(f"- imports: {len(imps) if isinstance(imps, list) else 0}")
            if isinstance(analysis, dict) and analysis:
                if analysis.get("summary"):
                    lines.append(f"- summary: {analysis.get('summary')}")
                rs = analysis.get("responsibilities") or []
                if isinstance(rs, list) and rs:
                    lines.append("- responsibilities:")
                    for r in rs[:10]:
                        lines.append(f"  - {r}")
                kis = analysis.get("key_interfaces") or []
                if isinstance(kis, list) and kis:
                    lines.append("- key_interfaces:")
                    for k in kis[:20]:
                        lines.append(f"  - {k}")
                if analysis.get("error"):
                    lines.append(f"- analysis_error: {analysis.get('error')}")
            lines.append("")

    os.makedirs(os.path.dirname(os.path.abspath(out_path)) or ".", exist_ok=True)
    with open(out_path, "w", encoding="utf-8") as f:
        f.write("\n".join(lines))


def run_fileatlas(cfg: FileAtlasConfig) -> ProjectMap:
    project = build_project_map(cfg)
    save_project_map(project, cfg.out_path)
    if cfg.txt_out_path:
        save_txt_report(project, cfg.txt_out_path)
    return project

