"""
Codebase Indexer untuk AI Search (Fase 2)
Crawl monorepo → chunk → embed via Ollama nomic-embed-text → simpan ke pgvector
"""

import os
import sys
import hashlib
import json
import time
import requests
import psycopg
from pathlib import Path
from dotenv import load_dotenv

load_dotenv()

PGVECTOR_HOST = os.getenv("PGVECTOR_HOST", "10.100.1.24")
PGVECTOR_PORT = int(os.getenv("PGVECTOR_PORT", "5433"))
PGVECTOR_DB   = os.getenv("PGVECTOR_DB", "geonet_project_search")
PGVECTOR_USER = os.getenv("PGVECTOR_USER", "geonet_ai")
PGVECTOR_PASS = os.getenv("PGVECTOR_PASSWORD", "")
OLLAMA_HOST   = os.getenv("OLLAMA_HOST", "http://10.100.1.14:11434")
EMBED_MODEL   = os.getenv("EMBED_MODEL", "nomic-embed-text")
REPO_ROOT     = Path(os.getenv("REPO_ROOT", r"D:\Project\app on git"))
CHUNK_SIZE    = int(os.getenv("CHUNK_SIZE", "400"))
CHUNK_OVERLAP = int(os.getenv("CHUNK_OVERLAP", "50"))

INCLUDE_EXTENSIONS = {
    ".php", ".ts", ".tsx", ".js", ".py",
    ".md", ".yaml", ".yml", ".sql", ".env.example",
    ".json",
}

EXCLUDE_DIRS = {
    "node_modules", "vendor", ".git", ".next", "__pycache__",
    "dist", "build", ".idea", ".vscode", "storage/logs",
    ".venv", "venv", "env", ".env",
}

EXCLUDE_FILES = {
    "package-lock.json", "composer.lock", "yarn.lock",
}

LANGUAGE_MAP = {
    ".php": "php", ".ts": "typescript", ".tsx": "tsx",
    ".js": "javascript", ".py": "python", ".md": "markdown",
    ".yaml": "yaml", ".yml": "yaml", ".sql": "sql",
    ".json": "json",
}


def get_repo_name(file_path: Path) -> str:
    try:
        rel = file_path.relative_to(REPO_ROOT)
        parts = rel.parts
        return parts[0] if len(parts) > 1 else "root"
    except ValueError:
        return "unknown"


def chunk_text(text: str, size: int = CHUNK_SIZE, overlap: int = CHUNK_OVERLAP):
    """Split teks jadi chunk berdasarkan baris, bukan token."""
    lines = text.splitlines(keepends=True)
    chunks = []
    current, current_lines, start_line = [], 0, 1

    for i, line in enumerate(lines, start=1):
        current.append(line)
        current_lines += 1
        if current_lines >= size:
            chunks.append(("".join(current), start_line, i))
            keep = max(0, len(current) - overlap)
            current = current[keep:]
            start_line = i - len(current) + 1
            current_lines = len(current)

    if current:
        chunks.append(("".join(current), start_line, start_line + len(current) - 1))

    return chunks


def get_embedding(text: str, retries: int = 3) -> list[float] | None:
    for attempt in range(retries):
        try:
            resp = requests.post(
                f"{OLLAMA_HOST}/api/embeddings",
                json={"model": EMBED_MODEL, "prompt": text},
                timeout=60,
            )
            resp.raise_for_status()
            return resp.json().get("embedding")
        except Exception as e:
            if attempt < retries - 1:
                time.sleep(2 ** attempt)  # exponential backoff: 1s, 2s
            else:
                print(f"  [EMBED ERROR] {e}")
    return None


def collect_files() -> list[Path]:
    files = []
    for path in REPO_ROOT.rglob("*"):
        if not path.is_file():
            continue
        if any(ex in path.parts for ex in EXCLUDE_DIRS):
            continue
        if path.name in EXCLUDE_FILES:
            continue
        if path.suffix.lower() not in INCLUDE_EXTENSIONS:
            continue
        files.append(path)
    return sorted(files)


def file_checksum(path: Path) -> str:
    return hashlib.md5(path.read_bytes()).hexdigest()


def index_files(conn, files: list[Path]):
    cur = conn.cursor()

    cur.execute("SELECT file_path, checksum FROM code_chunks GROUP BY file_path, checksum")
    indexed = {row[0]: row[1] for row in cur.fetchall()}

    total = len(files)
    skipped = inserted = errors = 0

    for i, fpath in enumerate(files, 1):
        rel_path = str(fpath.relative_to(REPO_ROOT)).replace("\\", "/")
        checksum = file_checksum(fpath)

        if indexed.get(rel_path) == checksum:
            skipped += 1
            continue

        try:
            text = fpath.read_text(encoding="utf-8", errors="ignore")
        except Exception as e:
            print(f"  [READ ERROR] {rel_path}: {e}")
            errors += 1
            continue

        if not text.strip():
            skipped += 1
            continue

        chunks = chunk_text(text)
        lang = LANGUAGE_MAP.get(fpath.suffix.lower(), "text")
        repo = get_repo_name(fpath)

        if rel_path in indexed:
            cur.execute("DELETE FROM code_chunks WHERE file_path = %s", (rel_path,))

        chunk_ok = 0
        for content, start, end in chunks:
            content = content.replace("\x00", "")
            if not content.strip():
                continue
            embedding = get_embedding(content)
            if embedding is None:
                continue
            cur.execute(
                """
                INSERT INTO code_chunks (repo, file_path, language, start_line, end_line, content, embedding, checksum)
                VALUES (%s, %s, %s, %s, %s, %s, %s::vector, %s)
                """,
                (repo, rel_path, lang, start, end, content, json.dumps(embedding), checksum),
            )
            chunk_ok += 1

        conn.commit()
        inserted += chunk_ok

        print(f"[{i}/{total}] {rel_path} → {chunk_ok} chunks", flush=True)

    cur.close()
    print(f"\n✅ Selesai: {inserted} chunks inserted, {skipped} files skipped, {errors} errors")


def main():
    print(f"Connecting to pgvector {PGVECTOR_HOST}:{PGVECTOR_PORT}...")
    conn = psycopg.connect(
        host=PGVECTOR_HOST, port=PGVECTOR_PORT, dbname=PGVECTOR_DB,
        user=PGVECTOR_USER, password=PGVECTOR_PASS,
    )
    print(f"Repo root: {REPO_ROOT}")
    print(f"Ollama embed: {OLLAMA_HOST} model={EMBED_MODEL}")

    files = collect_files()
    print(f"Files to index: {len(files)}")

    index_files(conn, files)
    conn.close()


if __name__ == "__main__":
    main()