modules/study_writer.py

import json
import os
import re
import time
from pathlib import Path

from config import (
    THESIS_DIR,
    OPENAI_MODEL,
    AI_FALLBACK_MODELS,
    RETRY_BASE_DELAY,
    SYSTEM_PROMPTS_DIR,
)
from modules.term_workspace import TermWorkspace
from modules.draft_builder import PromptBuilder
from modules.resilience import CircuitBreaker, circuit_breaker
from clients.openai_client import OpenAIClient

SYSTEM_PROMPTS = {
    "global": "00_global_rules.md",
    "tahdid": "01_tahdid.md",
    "lugha": "02_lugha.md",
    "tarikh": "03_tarikh.md",
    "madloul": "04_madloul.md",
    "muqarana": "05_muqarana.md",
    "khatima": "06_khatima.md",
    "masadir": "07_masadir.md",
    "reviewer": "08_reviewer.md",
}

SECTIONS = [
    ("tahdid", "التمهيد"),
    ("lugha", "المبحث الأول: الدراسة اللغوية"),
    ("tarikh", "المبحث الثاني: الدراسة الاصطلاحية"),
    ("madloul", "المبحث الثالث: استعمال القطان"),
    ("muqarana", "المبحث الرابع: الدراسة التطبيقية المقارنة"),
    ("khatima", "المبحث الخامس: الخاتمة"),
    ("masadir", "قائمة المصادر"),
]


class StudyWriter:
    def __init__(self, term: str, category: str, draft: dict, api_key: str = ""):
        self.term = term
        self.category = category
        self.draft = draft
        self.api_key = api_key
        self.ws = TermWorkspace(term, category)
        self.builder = PromptBuilder(term, category, draft)
        self.system_rules = self._load_system_prompt("global")
        self._client = OpenAIClient(api_key)

    @staticmethod
    def get_circuit_breaker_status() -> str:
        return circuit_breaker.status

    @staticmethod
    def estimate_tokens(text: str) -> int:
        """Rough token estimate: ~1 token per 3.5 chars for mixed Arabic/English."""
        if not text:
            return 0
        arabic_chars = sum(
            1 for c in text if "\u0600" <= c <= "\u06ff" or "\u0750" <= c <= "\u077f"
        )
        other_chars = len(text) - arabic_chars
        # Arabic: ~1.8 chars/token, English/code: ~4 chars/token
        return int(arabic_chars / 1.8 + other_chars / 4)

    @staticmethod
    def validate_prompt_sizes(
        system_content: str, user_content: str, max_input_tokens: int = 28000
    ) -> tuple[bool, str, int, int]:
        """Check if prompts fit within model context. Returns (ok, error_msg, sys_tokens, user_tokens)."""
        sys_tokens = StudyWriter.estimate_tokens(system_content)
        user_tokens = StudyWriter.estimate_tokens(user_content)
        total = sys_tokens + user_tokens
        if total > max_input_tokens:
            return (
                False,
                (
                    f"Input too large: ~{total:,} tokens (max {max_input_tokens:,}). "
                    f"System={sys_tokens:,} User={user_tokens:,}. "
                    f"Reduce raw material or split the section."
                ),
                sys_tokens,
                user_tokens,
            )
        return True, "", sys_tokens, user_tokens

    def _load_system_prompt(self, name: str) -> str:
        filename = SYSTEM_PROMPTS.get(name, "")
        if not filename:
            return ""
        path = SYSTEM_PROMPTS_DIR / filename
        if path.exists():
            return path.read_text(encoding="utf-8").format(term=self.term)
        return ""

    def _build_section_prompt(
        self, section_key: str, running_context: str = ""
    ) -> tuple[str, str]:
        section_specific = self._load_system_prompt(section_key)
        system_content = f"{self.system_rules}\n\n---\n\n{section_specific}"

        max_chars = {
            "tahdid": 5000,
            "lugha": 8000,
            "tarikh": 8000,
            "madloul": 12000,
            "muqarana": 12000,
            "khatima": 4000,
            "masadir": 8000,
        }.get(section_key, 8000)

        # A4: Section-specific raw material delivery
        section_draft = self.builder.build_section_draft(
            section_key, max_chars=max_chars
        )

        user_content = ""

        # Include analytical brief for first section to orient the AI
        if section_key == "tahdid":
            brief = self.builder.build_analytical_brief()
            user_content += f"## الملخص التحليلي — اقرأه أولاً لتعرف المادة المتاحة:\n{brief}\n\n---\n\n"

        if running_context:
            user_content += f"## سياق الأقسام السابقة (للمتابعة وتجنب التكرار):\n{running_context}\n\n---\n\n"

        user_content += f"""## المادة الخام — استخدمها كلها ولا تترك أي نتيجة:

{section_draft}

---

## ملخص الدراسة
- اللفظ: «{self.term}»
- التصنيف: {"تعديل" if self.category == "tadeel" else "جرح"}
- إجمالي النتائج: {self.draft.get("total_results", 0)}
- عدد الكتب: {self.draft.get("books_count", 0)}
- عدد المؤلفين: {self.draft.get("authors_count", 0)}

تنبيه حاسم: اكتب المحتوى النهائي مباشرة باللغة العربية الفصحى. ممنوع منعاً باتاً التفكير بصوت عالٍ أو كتابة أي مسودات أو تعليقات باللغة الإنجليزية. اكتب المبحث الآن."""

        # Safety net: if the raw material still exceeds the context budget,
        # truncate the material block instead of failing hard.
        max_user_tokens = 24000
        usr_tok = self.estimate_tokens(user_content)
        sys_tok = self.estimate_tokens(system_content)
        if sys_tok + usr_tok > max_user_tokens:
            allowed = max(2000, max_user_tokens - sys_tok)
            allowed_chars = int(allowed * 2.2)  # Arabic ≈ 1.8–2.2 chars/token
            truncated_draft = section_draft[:allowed_chars]
            if len(truncated_draft) < len(section_draft):
                truncated_draft += (
                    "\n\n> ⚠️ **تنبيه**: المادة الخام اقتُطعت لتجاوز حد السياق. "
                    "النتائج الكاملة في draft.md."
                )
            user_content = (
                user_content[: user_content.find("## المادة الخام")]
                + f"## المادة الخام — استخدمها كلها ولا تترك أي نتيجة:\n\n{truncated_draft}"
                + user_content[user_content.find("---\n\n## ملخص الدراسة") :]
            )
            self._log(
                f"Auto-truncated raw material for {section_key}: "
                f"~{usr_tok:,} -> ~{self.estimate_tokens(user_content):,} user tokens"
            )
        return system_content, user_content

    def write(self, progress_callback=None, resume: bool = True) -> tuple[str, Path]:
        self._client = OpenAIClient(self.api_key)
        if not self._client.api_key:
            raise ValueError("OPENAI_API_KEY غير موجود في البيئة")

        self._report(progress_callback, "جاري تجهيز المواد الخام...", 5)

        md_draft = self.builder.build_markdown_draft()

        self.ws.save_raw_draft_md(md_draft)
        self.ws.save_research_before_ai(md_draft)
        self.ws.save_raw_search(self.draft)
        self.ws.save_meta(
            {
                "status": "writing_started",
                "total_results": self.draft.get("total_results", 0),
            }
        )
        self._log(
            f"Starting study writing. Total results: {self.draft.get('total_results', 0)}"
        )

        self._report(progress_callback, f"تم حفظ المواد الخام في {self.ws.raw_dir}", 8)

        all_parts = []
        running_context = ""
        for i, (key, section_name) in enumerate(SECTIONS):
            pct_start = 10 + (i * 12)
            pct_end = 10 + ((i + 1) * 12)

            if resume:
                ckpt_path = self.ws.logs_dir / f"checkpoint_{key}.json"
                if ckpt_path.exists():
                    try:
                        ckpt = json.loads(ckpt_path.read_text(encoding="utf-8"))
                        if ckpt.get("status") == "completed" and ckpt.get("content"):
                            part_text = ckpt["content"]
                            self._report(
                                progress_callback,
                                f"استئناف: {section_name} موجود مسبقاً",
                                pct_end,
                            )
                            self.ws.save_ai_section(i + 1, key, part_text)
                            all_parts.append(f"## {section_name}\n\n{part_text}")
                            max_ctx = 2000
                            truncated = part_text[:max_ctx]
                            running_context += f"\n## {section_name}:\n{truncated}\n"
                            continue
                    except Exception as e:
                        self._log(f"Failed to read checkpoint for {key}: {e}")

            self._report(
                progress_callback,
                f"كتابة: {section_name} ({i + 1}/{len(SECTIONS)})...",
                pct_start,
            )

            system_prompt, user_prompt = self._build_section_prompt(
                key, running_context
            )
            self.ws.save_prompt(
                f"# SYSTEM PROMPT:\n{system_prompt}\n\n---\n\n# USER PROMPT:\n{user_prompt}",
                key,
            )

            self._save_checkpoint(key, "started")
            part_text = self._call_with_retry(
                self._client.api_key,
                system_prompt,
                user_prompt,
                progress_callback,
                pct_start,
                pct_end,
            )

            self.ws.save_ai_section(i + 1, key, part_text)
            self._save_checkpoint(key, "completed", part_text)
            all_parts.append(f"## {section_name}\n\n{part_text}")

            # Pass full previous section for continuity (capped to avoid prompt explosion)
            max_context_per_section = 2000
            truncated = part_text[:max_context_per_section]
            if len(part_text) > max_context_per_section:
                truncated += "\n[...تم اقتطاع السياق...] "
            running_context += f"\n## {section_name}:\n{truncated}\n"

            if i < len(SECTIONS) - 1:
                self._report(
                    progress_callback,
                    "انتظار 25 ثانية لتجنب حدود معدل الطلبات...",
                    pct_end,
                )
                time.sleep(25)

        full_text = "\n\n".join(all_parts)

        self.ws.save_ai_full(full_text)
        self._report(progress_callback, "جاري التنسيق...", 92)
        full_text = self._post_process(full_text)

        self._report(progress_callback, "جاري الحفظ...", 95)
        file_path = self._save_study(full_text)
        self.ws.save_processed(full_text)
        self.ws.save_meta({"status": "writing_completed", "file": str(file_path)})

        self._report(progress_callback, f"تم الحفظ في {file_path.name}", 100)

        return full_text, file_path

    def _call_with_retry(
        self,
        api_key,
        system_content,
        user_content,
        progress_callback,
        pct_start,
        pct_end,
    ):
        models = []
        for m in [OPENAI_MODEL] + AI_FALLBACK_MODELS:
            if m not in models:
                models.append(m)
        last_error = None

        # Validate prompt sizes before calling
        ok, err_msg, sys_t, usr_t = self.validate_prompt_sizes(
            system_content, user_content
        )
        if not ok:
            self._report(progress_callback, f"تنبيه: {err_msg}", pct_start)
            self._log(
                f"PROMPT_TOO_LARGE: sys≈{sys_t} tok, user≈{usr_t} tok — {err_msg}"
            )
            raise ValueError(err_msg)

        self._log(
            f"Prompt size: system≈{sys_t:,} tok, user≈{usr_t:,} tok, "
            f"total≈{sys_t + usr_t:,} tok, circuit_breaker={circuit_breaker.status}"
        )

        # Check circuit breaker
        if not circuit_breaker.allow_request():
            self._report(
                progress_callback,
                f"Circuit breaker OPEN — API فاشل متكرراً، انتظار...",
                pct_start,
            )
            raise Exception(
                f"Circuit breaker is OPEN ({circuit_breaker.status}). "
                f"API failing consecutively. Try again later."
            )

        for attempt in range(3):
            for model_idx, model in enumerate(models):
                try:
                    self._report(
                        progress_callback, f"جاري الكتابة بـ {model}...", pct_start + 3
                    )
                    # Try streaming first, fall back to non-streaming
                    try:
                        tokens = list(
                            self._client.call_stream(
                                model, system_content, user_content
                            )
                        )
                        text = self._client.last_stream_result
                        if text:
                            self._report(
                                progress_callback,
                                f"تم ({len(text)} حرف، {len(tokens)} توكن)",
                                pct_end,
                            )
                            circuit_breaker.record_success()
                            return text
                    except Exception as stream_err:
                        self._log(
                            f"Stream failed ({stream_err}), falling back to non-streaming"
                        )
                        text = self._client.call(model, system_content, user_content)

                    if not text or len(text) < 100:
                        raise ValueError(f"الرد قصير جداً ({len(text or '')} حرف)")
                    self._report(progress_callback, f"تم ({len(text)} حرف)", pct_end)
                    circuit_breaker.record_success()
                    return text
                except Exception as e:
                    last_error = e
                    circuit_breaker.record_failure()
                    err_str = str(e)
                    self._report(progress_callback, f"فشل: {err_str[:80]}", pct_start)

                    # 402 = insufficient credits — no point retrying other models
                    if "402" in err_str:
                        self._report(
                            progress_callback,
                            f"خطأ 402: رصيد غير كافي في {model} — جرب مفتاح API آخر",
                            pct_start,
                        )
                        raise ValueError(
                            f"رصيد OpenRouter غير كافي. "
                            f"أضف رصيداً على https://openrouter.ai/settings/credits "
                            f"أو أدخل مفتاح API جديد في حقل 'مفتاح OpenRouter API'."
                        )

                    if model_idx < len(models) - 1:
                        sleep_time = 45 if "429" in err_str else 5
                        self._report(
                            progress_callback,
                            f"فشل في {model}، انتظار {sleep_time}s ثم تجربة {models[model_idx + 1]}",
                            pct_start,
                        )
                        time.sleep(sleep_time)
                        continue
                    else:
                        delay = max(45, RETRY_BASE_DELAY * (2**attempt))
                        self._report(
                            progress_callback,
                            f"فشلت جميع النماذج، انتظار {delay}s قبل المحاولة القادمة",
                            pct_start,
                        )
                        time.sleep(delay)
                        break
        raise last_error or Exception("فشلت جميع المحاولات")

    def _post_process(self, text: str) -> str:
        text = re.sub(r"\n{3,}", "\n\n", text)
        text = re.sub(r"`(.*?)`", r"«\1»", text)
        return text.strip()

    def review(self, text: str, progress_callback=None) -> str:
        api_key = self.api_key
        model = OPENAI_MODEL

        review_specific = self._load_system_prompt("reviewer")
        system_content = f"{self.system_rules}\n\n---\n\n{review_specific}"
        user_content = f"النص الأكاديمي المطلوب مراجعته وتدقيقه لغوياً ومنهجياً وتنسيقياً:\n---\n{text}\n---"

        self._report(progress_callback, "جاري المراجعة اللغوية...", 10)

        for attempt in range(2):
            try:
                self._report(progress_callback, f"جاري المراجعة بـ {model}...", 30)
                corrected = self._client.call(model, system_content, user_content)
                if not corrected or len(corrected) < len(text) * 0.5:
                    raise ValueError(f"الرد قصير جداً ({len(corrected)} حرف)")
                self._report(progress_callback, "تمت المراجعة", 100)
                return corrected
            except Exception as e:
                err_str = str(e)
                for key_var in ["OPENAI_API_KEY", "OPENROUTER_API_KEY"]:
                    val = os.environ.get(key_var, "")
                    if val:
                        err_str = err_str.replace(val, "***")
                self._report(progress_callback, f"فشل: {err_str[:50]}", 10)
                if "429" in err_str:
                    time.sleep(RETRY_BASE_DELAY * (attempt + 1))
                else:
                    time.sleep(2)

        raise Exception("فشلت المراجعة")

    def _save_study(self, text: str) -> Path:
        # Primary: save inside workspace processed dir
        primary_path = self.ws.processed_dir / "study.md"
        primary_path.write_text(text, encoding="utf-8")

        # Secondary: append to master thesis file (Rule 8: no side files)
        master_path = (
            THESIS_DIR / "رسالة_ألفاظ_الجرح_والتعديل_عند_يحيى_بن_سعيد_القطان.md"
        )
        try:
            separator = f"\n\n---\n\n## {self.term}\n\n"
            if master_path.exists():
                existing = master_path.read_text(encoding="utf-8")
                if self.term not in existing:
                    master_path.write_text(
                        existing + separator + text, encoding="utf-8"
                    )
            else:
                master_path.write_text(
                    f"# رسالة ألفاظ الجرح والتعديل عند يحيى بن سعيد القطان (ت 198 هـ)\n\n"
                    + separator
                    + text,
                    encoding="utf-8",
                )
        except Exception as e:
            self._log(f"Failed to append to master thesis: {e}")

        # Update master index
        index_path = THESIS_DIR / "01_terms" / "all-terms.md"
        if index_path.exists():
            try:
                content = index_path.read_text(encoding="utf-8")
                if self.term not in content:
                    with index_path.open("a", encoding="utf-8") as f:
                        f.write(f"\n- [مكتمل] {self.term}")
            except Exception:
                pass

        self._log(f"Study saved: primary={primary_path} master={master_path}")
        return primary_path

    def _report(self, callback, message: str, progress: int):
        self._log(f"Report: progress={progress}% message={message}")
        if callback:
            callback(message, progress)

    def _log(self, msg: str):
        log_path = self.ws.logs_dir / "session.log"
        timestamp = time.strftime("%Y-%m-%d %H:%M:%S")
        try:
            with log_path.open("a", encoding="utf-8") as f:
                f.write(f"[{timestamp}] {msg}\n")
        except Exception:
            pass

    def rewrite_section(
        self, section_key: str, section_name: str, full_context: str, api_key: str = ""
    ) -> str:
        """Regenerate a single section with full context from all other sections."""
        api_key = api_key or self.api_key
        if not api_key:
            raise ValueError("OPENAI_API_KEY غير موجود في البيئة")

        system_prompt, user_prompt = self._build_section_prompt(
            section_key, full_context
        )

        # Add instruction to rewrite
        user_prompt = (
            f"## مهمة إضافية: إعادة كتابة هذا القسم\n"
            f"تم اكتشاف جودة منخفضة في هذا القسم. أعد كتابته بالكامل مع التزام تام "
            f"بوضع توثيق مباشر في كل فقرة.\n\n"
            f"---\n\n{user_prompt}"
        )

        self.ws.save_prompt(
            f"# REWRITE SYSTEM:\n{system_prompt}\n\n---\n\n# REWRITE USER:\n{user_prompt}",
            f"{section_key}_rewrite",
        )

        text = self._call_with_retry(api_key, system_prompt, user_prompt, None, 0, 100)
        self.ws.save_ai_section(0, f"{section_key}_rewrite", text)
        self._save_checkpoint(section_key, "rewritten", text)
        return text

    def _save_checkpoint(self, section_key: str, status: str, content: str = ""):
        word_count = len(content.split()) if content else 0
        citation_count = (
            len(re.findall(r"🔗|https?://|يحتاج توثيقاً", content)) if content else 0
        )
        checkpoint_data = {
            "section": section_key,
            "status": status,
            "content": content,
            "word_count": word_count,
            "citation_count": citation_count,
            "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
        }
        try:
            path = self.ws.logs_dir / f"checkpoint_{section_key}.json"
            path.write_text(
                json.dumps(checkpoint_data, ensure_ascii=False, indent=2),
                encoding="utf-8",
            )
            self._log(
                f"Checkpoint saved for {section_key}: status={status} words={word_count} citations={citation_count}"
            )
        except Exception as e:
            self._log(f"Failed to save checkpoint for {section_key}: {e}")