#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
passport_auto_enricher.py
=========================
Автоматический обогатитель единого паспорта знака «Оракула 999».
Выполняет последовательный прогон 78 атомарных промптов на каждый символ (000–999).
Сохраняет состояние после каждого шага (checkpoint-per-field).
"""

import os
import sys
import json
import time
import argparse
import urllib.request
import urllib.error
from pathlib import Path
from typing import Dict, Any, List, Optional

# Add project root to path
PROJECT_ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(PROJECT_ROOT))

from oracle_enricher.enricher.passport_schema import (
    PROMPTS_78,
    PROMPT_BY_KEY,
    PROMPT_BY_INDEX,
    PassportPromptDef,
    render_passport_markdown,
    count_filled_fields,
    empty_passport_record,
    get_field_prompt,
    get_all_78_prompts,
)

DB_PATH = PROJECT_ROOT / "assets" / "data" / "oracle_passports.json"
FALLBACK_DB_PATH = PROJECT_ROOT / "oracle_enricher" / "data" / "oracle_passports.json"
OUT_PASSPORT_DIR = PROJECT_ROOT / "assets" / "passports"


def load_database() -> Dict[str, Dict[str, Any]]:
    path = DB_PATH if DB_PATH.exists() else FALLBACK_DB_PATH
    if not path.exists():
        print(f"[ERROR] Database file not found at {path}")
        sys.exit(1)
    with open(path, "r", encoding="utf-8") as f:
        return json.load(f)


def save_database(data: Dict[str, Dict[str, Any]]) -> None:
    # Atomic write to assets/data/oracle_passports.json
    DB_PATH.parent.mkdir(parents=True, exist_ok=True)
    temp_path = DB_PATH.with_suffix(".tmp")
    with open(temp_path, "w", encoding="utf-8") as f:
        json.dump(data, f, ensure_ascii=False, indent=2)
    os.replace(temp_path, DB_PATH)

    # Sync to oracle_enricher/data if folder exists
    if FALLBACK_DB_PATH.parent.exists():
        try:
            with open(FALLBACK_DB_PATH, "w", encoding="utf-8") as f:
                json.dump(data, f, ensure_ascii=False, indent=2)
        except Exception:
            pass


class LLMClient:
    def __init__(
        self,
        backend: str = "dry-run",
        model: str = "gpt-4o-mini",
        api_key: Optional[str] = None,
        base_url: Optional[str] = None,
        timeout: int = 45,
    ):
        self.backend = backend.lower()
        self.model = model
        self.api_key = api_key or os.environ.get("OPENAI_API_KEY") or os.environ.get("OPENROUTER_API_KEY") or os.environ.get("DEEPSEEK_API_KEY") or ""
        self.base_url = (base_url or os.environ.get("OPENAI_BASE_URL") or "https://api.openai.com/v1").rstrip("/")
        self.timeout = timeout

    def generate(self, prompt: str, system_prompt: str = "") -> str:
        if self.backend == "dry-run":
            # Simulate or return dry run indication
            return f"[DRY-RUN] {prompt[:60]}..."

        elif self.backend == "mock":
            # Return plausible mock responses
            return "[Каноническое заполнение поля по заданному формату]"

        elif self.backend in ("openai", "openrouter", "deepseek", "custom"):
            return self._call_openai_compatible(prompt, system_prompt)

        elif self.backend == "ollama":
            return self._call_ollama(prompt, system_prompt)

        else:
            raise ValueError(f"Unknown backend: {self.backend}")

    def _call_openai_compatible(self, prompt: str, system_prompt: str) -> str:
        url = f"{self.base_url}/chat/completions"
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.api_key}",
        }
        messages = []
        if system_prompt:
            messages.append({"role": "system", "content": system_prompt})
        messages.append({"role": "user", "content": prompt})

        payload = {
            "model": self.model,
            "messages": messages,
            "temperature": 0.4,
            "max_tokens": 500,
        }

        req = urllib.request.Request(url, data=json.dumps(payload).encode("utf-8"), headers=headers, method="POST")
        with urllib.request.urlopen(req, timeout=self.timeout) as resp:
            body = json.loads(resp.read().decode("utf-8"))
            return body["choices"][0]["message"]["content"].strip()

    def _call_ollama(self, prompt: str, system_prompt: str) -> str:
        url = f"{self.base_url}/api/generate" if "11434" in self.base_url else "http://localhost:11434/api/generate"
        payload = {
            "model": self.model,
            "prompt": prompt,
            "system": system_prompt,
            "stream": False,
            "options": {"temperature": 0.4},
        }
        headers = {"Content-Type": "application/json"}
        req = urllib.request.Request(url, data=json.dumps(payload).encode("utf-8"), headers=headers, method="POST")
        with urllib.request.urlopen(req, timeout=self.timeout) as resp:
            body = json.loads(resp.read().decode("utf-8"))
            return body.get("response", "").strip()


def clean_llm_response(text: str, expected_key: str) -> str:
    """Clean extra fluff, quotes, or markdown code blocks from LLM output."""
    t = text.strip()
    if t.startswith("```") and t.endswith("```"):
        lines = t.splitlines()
        if len(lines) >= 3:
            t = "\n".join(lines[1:-1]).strip()
    if (t.startswith('"') and t.endswith('"')) or (t.startswith("'") and t.endswith("'")):
        t = t[1:-1].strip()
    return t


def enrich_symbol(
    entry: Dict[str, Any],
    client: LLMClient,
    db: Dict[str, Dict[str, Any]],
    force: bool = False,
    section_filter: Optional[str] = None,
    field_filter: Optional[str] = None,
    delay: float = 0.3,
    max_retries: int = 3,
) -> int:
    idx_num = entry.get("id", 0)
    idx_str = f"{idx_num:03d}"
    word = entry.get("word", "")
    emoji = entry.get("emoji", "")
    fields = entry.setdefault("fields", {})

    print(f"\n======================================================================")
    print(f"🔮 НАЧАЛО ОБОГАЩЕНИЯ: [{idx_str}] «{word.upper()}» {emoji}")
    print(f"   Заполнено: {entry.get('filled_count', 0)}/78 | Статус: {'ГОТОВО' if entry.get('is_complete') else 'В ПРОЦЕССЕ'}")
    print(f"======================================================================")

    enriched_this_run = 0

    system_prompt = (
        "Ты — сакральный историограф, лингвист, герменевт и архивариус системы «Оракул 999».\n"
        "Твоя задача — формировать канонический единый паспорт знака строго по заданному формату.\n"
        "Пиши предельно емко, глубоко, без лишних предисловий и вежливостей, строго соблюдая указанный шаблон."
    )

    for p_def in PROMPTS_78:
        # Filter checks
        if section_filter and p_def.section_num != section_filter:
            continue
        if field_filter and field_filter not in (p_def.key, str(p_def.index)):
            continue

        curr_val = fields.get(p_def.key, "")
        if curr_val and not force:
            continue

        prompt_text = get_field_prompt(p_def.key, entry)

        print(f"  [{idx_str}|{word}] [{p_def.index:02d}/78] Разд.{p_def.section_num} • {p_def.title}...", end="", flush=True)

        resp_text = ""
        duration = 0.0
        success = False
        for attempt in range(max_retries):
            try:
                t0 = time.time()
                raw_resp = client.generate(prompt_text, system_prompt=system_prompt)
                resp_text = clean_llm_response(raw_resp, p_def.key)
                duration = time.time() - t0
                success = True
                break
            except Exception as e:
                print(f" [Попытка {attempt+1} сбой: {e}]", end="", flush=True)
                time.sleep(2.0 * (attempt + 1))

        if not success:
            print(" ❌ ОШИБКА")
            continue

        fields[p_def.key] = resp_text
        filled, _ = count_filled_fields(entry)
        entry["filled_count"] = filled
        entry["is_complete"] = (filled >= 78)
        entry["updated_at"] = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())

        # Real-time checkpoint save to disk!
        if client.backend != "dry-run":
            save_database(db)
        enriched_this_run += 1

        short_preview = (resp_text[:50] + "...") if len(resp_text) > 50 else resp_text
        print(f" ✅ ({duration:.1f}s) -> {short_preview}")

        if delay > 0:
            time.sleep(delay)

    # Export markdown if complete or substantial
    if entry.get("filled_count", 0) >= 70:
        OUT_PASSPORT_DIR.mkdir(parents=True, exist_ok=True)
        md_file = OUT_PASSPORT_DIR / f"{idx_str}_{word}.md"
        with open(md_file, "w", encoding="utf-8") as f:
            f.write(render_passport_markdown(entry))
        print(f"  📄 Сохранен паспорт в Markdown: {md_file.name}")

    print(f"🏁 Завершено для [{idx_str} «{word}»]: обновлено {enriched_this_run} полей (всего {entry.get('filled_count', 0)}/78).")
    return enriched_this_run


def print_stats(db: Dict[str, Dict[str, Any]]) -> None:
    total = len(db)
    filled_counts = [entry.get("filled_count", 0) for entry in db.values()]
    complete_count = sum(1 for c in filled_counts if c >= 78)
    partially_filled = sum(1 for c in filled_counts if c > 0)
    total_fields = sum(filled_counts)
    max_possible = total * 78

    print("\n" + "=" * 65)
    print("📊 СТАТИСТИКА БАЗЫ ДАННЫХ «ОРАКУЛ 999 — ПАСПОРТА ЗНАКОВ»")
    print("=" * 65)
    print(f" Всего знаков в реестре:     {total}")
    print(f" Полностью завершено (78/78): {complete_count} ({complete_count/total*100:.1f}%)")
    print(f" Частично заполнено (>0):    {partially_filled} ({partially_filled/total*100:.1f}%)")
    print(f" Всего заполненных ячеек:    {total_fields} / {max_possible} ({total_fields/max_possible*100:.2f}%)")
    print(f" Среднее кол-во полей/карту: {total_fields/total:.1f} из 78")

    # Section breakdown
    section_counts: Dict[str, int] = {f"Раздел {s}": 0 for s in ["I", "II", "III", "IV", "V", "VI"]}
    for entry in db.values():
        flds = entry.get("fields", {})
        for p in PROMPTS_78:
            if flds.get(p.key):
                section_counts[f"Раздел {p.section_num}"] += 1

    print("\nРаспределение заполненности по разделам:")
    for sec, cnt in section_counts.items():
        sec_total = sum(1 for p in PROMPTS_78 if f"Раздел {p.section_num}" == sec) * total
        print(f"  {sec:<12}: {cnt:>6} / {sec_total} ({cnt/sec_total*100:.1f}%)")
    print("=" * 65 + "\n")


def export_prompts_file(entry: Dict[str, Any], out_path: Path) -> None:
    idx_num = entry.get("id", 0)
    idx_str = f"{idx_num:03d}"
    word = entry.get("word", "")
    emoji = entry.get("emoji", "")

    lines = [
        f"# ПАКЕТ 78 ПРОМПТОВ ДЛЯ ЗНАКА: [{idx_str}] «{word.upper()}» {emoji}",
        f"# Сгенерировано: {time.strftime('%Y-%m-%d %H:%M:%S')}",
        "=" * 70,
        "",
    ]
    for p in PROMPTS_78:
        p_text = get_field_prompt(p.key, entry)
        lines.append(f"### [ПРОМПТ {p.index:02d}/78] | Раздел {p.section_num} | Поле: {p.key}")
        lines.append(f"# Название: {p.title}")
        lines.append(p_text)
        lines.append("")
        lines.append("-" * 70)
        lines.append("")

    out_path.parent.mkdir(parents=True, exist_ok=True)
    with open(out_path, "w", encoding="utf-8") as f:
        f.write("\n".join(lines))
    print(f"✅ Экспортировано 78 промптов в файл: {out_path}")


def main():
    parser = argparse.ArgumentParser(
        description="Автоматический обогатитель единого паспорта знака «Оракула 999» (78 атомарных промптов)"
    )
    parser.add_argument("--stats", action="store_true", help="Показать общую статистику заполненности базы")
    parser.add_argument("--id", type=str, help="ID знака (например '000' или '0,1,2')")
    parser.add_argument("--word", type=str, help="Слово для обогащения (например 'огонь')")
    parser.add_argument("--range", nargs=2, type=int, metavar=("START", "END"), help="Диапазон индексов (например 0 10)")
    parser.add_argument("--all", action="store_true", help="Пройти по всем 1000 знакам")
    parser.add_argument("--backend", default="dry-run", choices=["dry-run", "mock", "openai", "openrouter", "deepseek", "ollama", "custom"], help="Бэкенд генерации")
    parser.add_argument("--model", default="gpt-4o-mini", help="Имя модели LLM")
    parser.add_argument("--api-key", default=None, help="API ключ")
    parser.add_argument("--base-url", default=None, help="Базовый URL OpenAI/Ollama")
    parser.add_argument("--section", choices=["I", "II", "III", "IV", "V", "VI", "1", "2", "3", "4", "5", "6"], help="Ограничить определенным разделом")
    parser.add_argument("--field", help="Ограничить конкретным ключом поля или индексом (1..78)")
    parser.add_argument("--force", action="store_true", help="Перезаписывать уже заполненные поля")
    parser.add_argument("--delay", type=float, default=0.2, help="Задержка между запросами (сек)")
    parser.add_argument("--export-prompts", action="store_true", help="Только экспортировать 78 промптов в файл out/prompts/...")

    args = parser.parse_args()

    db = load_database()

    if args.stats:
        print_stats(db)
        return

    sec_map = {"1": "I", "2": "II", "3": "III", "4": "IV", "5": "V", "6": "VI"}
    section_filter = sec_map.get(args.section, args.section)

    target_entries: List[Dict[str, Any]] = []

    if args.id:
        ids = [x.strip() for x in args.id.split(",")]
        for i_str in ids:
            int_val = int(i_str)
            k = str(int_val)
            if k in db:
                target_entries.append(db[k])
            else:
                print(f"[WARN] ID {i_str} не найден в базе!")

    elif args.word:
        w_lower = args.word.strip().lower()
        matched = [entry for entry in db.values() if entry.get("word", "").lower() == w_lower]
        if matched:
            target_entries.extend(matched)
        else:
            print(f"[WARN] Слово «{args.word}» не найдено в базе!")

    elif args.range:
        st, ed = args.range
        for i in range(st, ed + 1):
            k = str(i)
            if k in db:
                target_entries.append(db[k])

    elif args.all:
        for i in range(1000):
            k = str(i)
            if k in db:
                target_entries.append(db[k])

    else:
        print("💡 Не указан знак для обогащения. Используйте:")
        print("   --stats                  Показать статистику")
        print("   --word огонь             Обогатить по слову")
        print("   --id 000                 Обогатить по индексу")
        print("   --range 0 5              Обогатить диапазон")
        print("   --backend dry-run        Тестовый запуск без реальных API вызовов")
        print("   --export-prompts         Сформировать файл с 78 промптами")
        print_stats(db)
        return

    if not target_entries:
        print("⚠️ Целевые знаки не найдены.")
        return

    if args.export_prompts:
        for entry in target_entries:
            idx_str = f"{entry.get('id', 0):03d}"
            word = entry.get("word", "")
            out_file = PROJECT_ROOT / "out" / "prompts" / f"passport_{idx_str}_{word}_78prompts.txt"
            export_prompts_file(entry, out_file)
        return

    client = LLMClient(
        backend=args.backend,
        model=args.model,
        api_key=args.api_key,
        base_url=args.base_url,
    )

    print(f"🚀 Запуск обогатителя для {len(target_entries)} знаков.")
    print(f"   Бэкенд: {client.backend} | Модель: {client.model}")
    print(f"   Режим перезаписи: {'Да (--force)' if args.force else 'Только пустые ячейки'}")

    total_enriched = 0
    t_start = time.time()
    for entry in target_entries:
        count = enrich_symbol(
            entry=entry,
            client=client,
            db=db,
            force=args.force,
            section_filter=section_filter,
            field_filter=args.field,
            delay=args.delay,
        )
        total_enriched += count

    duration = time.time() - t_start
    print(f"\n🎉 ВСЕ ЗАДАЧИ ВЫПОЛНЕНЫ ЗА {duration:.1f}с!")
    print(f"   Всего обогащено полей: {total_enriched}")
    print_stats(db)


if __name__ == "__main__":
    main()