src/correction_service.py

import json
import re
import uuid
from datetime import datetime, timezone
from typing import Dict, List, Optional, Tuple

from .knowledge_base import KnowledgeBase
from .llm_client import GeminiClient
from .correction_store import CorrectionStore
from .models import (
    CorrectionRequest,
    CorrectionResponse,
    CorrectionRecord,
    RootCause,
)

# ---------------------------------------------------------------------------
# Prompts
# ---------------------------------------------------------------------------

_ROOT_CAUSE_PROMPT = """You are a RAG system quality analyst for a kidney stone clinical reporting system.

TASK: Determine why the AI-generated clinical report contained an error.

=== ORIGINAL REPORT SECTION: {section} ===
{original_section_content}

=== DOCTOR'S CORRECTION ===
{correction_text}

=== RETRIEVED CHUNKS USED TO GENERATE THE REPORT ===
{retrieved_chunks}

=== STONE FEATURES (ground truth) ===
{features_json}

Analyze carefully and classify the root cause as ONE of:

1. CHUNK_ERROR — A retrieved chunk contained wrong, outdated, or misleading medical information
   that caused the LLM to generate an incorrect statement. Identify which chunk(s) are problematic.

2. RETRIEVAL_MISS — The correct information EXISTS in the guidelines but was NOT retrieved.
   The retrieved chunks do not cover the corrected topic at all.

3. LLM_HALLUCINATION — The LLM generated information not supported by ANY retrieved chunk
   AND not derivable from the stone features. The retrieval was fine but the LLM fabricated details.

Respond in JSON only:
{{
  "root_cause": "CHUNK_ERROR" | "RETRIEVAL_MISS" | "LLM_HALLUCINATION",
  "reasoning": "<2-3 sentences explaining why>",
  "problematic_chunk_indices": [<0-based indices into retrieved_chunks list, or empty>],
  "new_chunk_suggestion": "<text to add as new chunk if RETRIEVAL_MISS, else null>",
  "new_chunk_metadata_hints": {{
    "content_type": "recommendation|narrative|table",
    "topics": "<comma separated>",
    "stone_size_range": "<1cm|1-2cm|>2cm or empty>",
    "stone_location": "<location or empty>",
    "evidence_level": "<level or empty>",
    "guideline_source": "EAU|AUA|Unknown"
  }}
}}"""

_CHUNK_FIX_PROMPT = """You are a medical knowledge base editor. A doctor has identified that
a clinical guideline chunk in our database contained wrong or misleading information.

=== ORIGINAL CHUNK CONTENT ===
{original_content}

=== DOCTOR'S CORRECTION ===
{correction_text}

=== CORRECTED SNIPPET (if provided) ===
{corrected_snippet}

Rewrite the chunk to fix the error while:
1. Preserving the medical terminology and formal style
2. Keeping the chunk roughly the same length (within 20%)
3. Retaining all correct medical information from the original
4. Incorporating the doctor's correction accurately

Output ONLY the corrected chunk text, no preamble."""


# ---------------------------------------------------------------------------
# Service
# ---------------------------------------------------------------------------

class CorrectionService:
    def __init__(
        self,
        kb: KnowledgeBase,
        llm: Optional[GeminiClient] = None,
        store: Optional[CorrectionStore] = None,
    ):
        self.kb = kb
        self.llm = llm or GeminiClient()
        self.store = store or CorrectionStore()

    # -----------------------------------------------------------------------
    # Public
    # -----------------------------------------------------------------------

    def process_correction(self, req: CorrectionRequest) -> CorrectionResponse:
        correction_id = str(uuid.uuid4())
        timestamp = datetime.now(timezone.utc).isoformat()

        root_cause, analysis = self._classify_root_cause(req)
        action_taken, affected_ids = self._remediate(req, root_cause, analysis, correction_id)

        # Audit kaydını correction_id ile güncelle (chunk_edits içindeki PENDING ID'leri düzelt)
        self._fix_pending_edit_ids(correction_id)

        response = CorrectionResponse(
            correction_id=correction_id,
            root_cause=root_cause,
            action_taken=action_taken,
            affected_chunk_ids=affected_ids,
            verification_query=self._build_verification_query(req),
        )

        record = CorrectionRecord(
            correction_id=correction_id,
            session_id=req.session_id,
            timestamp=timestamp,
            section=req.section,
            correction_text=req.correction_text,
            corrected_snippet=req.corrected_snippet,
            root_cause=root_cause,
            action_taken=action_taken,
            affected_chunk_ids=json.dumps(affected_ids),
            original_report=req.original_report,
            features_json=json.dumps(req.features),
            retrieved_context_json=json.dumps(
                [c.model_dump() for c in req.retrieved_context]
            ),
        )
        self.store.save_correction(record)

        return response

    def verify_correction(self, correction_id: str, features: Dict) -> Dict:
        record = self.store.get_correction(correction_id)
        if not record:
            raise ValueError(f"Correction {correction_id} not found")

        from .query_builder import build_all_queries

        queries = build_all_queries(features)
        largest_mm = features.get("largest_stone_mm")
        filter_dict = None
        if largest_mm:
            if largest_mm < 10:
                size_filter = "<1cm"
            elif largest_mm <= 20:
                size_filter = "1-2cm"
            else:
                size_filter = ">2cm"
            filter_dict = {
                "$or": [
                    {"stone_size_range": size_filter},
                    {"stone_size_range": ""},
                ]
            }

        seen: set = set()
        new_results = []
        for q in queries:
            for r in self.kb.search_with_chroma_ids(q, k=3, filter_dict=filter_dict):
                key = r["content"][:100]
                if key not in seen:
                    seen.add(key)
                    new_results.append(r)

        affected_ids = json.loads(record["affected_chunk_ids"])
        retrieved_ids = [r.get("chroma_id") for r in new_results]

        summary = {
            "correction_id": correction_id,
            "root_cause": record["root_cause"],
            "affected_chunk_ids": affected_ids,
            "new_retrieval_count": len(new_results),
            "affected_chunks_still_retrieved": [
                cid for cid in affected_ids if cid in retrieved_ids
            ],
        }
        self.store.mark_verified(correction_id, json.dumps(summary))
        return summary

    # -----------------------------------------------------------------------
    # Root cause analysis
    # -----------------------------------------------------------------------

    def _classify_root_cause(
        self, req: CorrectionRequest
    ) -> Tuple[RootCause, Dict]:
        chunks_text = ""
        for i, chunk in enumerate(req.retrieved_context):
            chunks_text += (
                f"\n[Chunk {i}] Source: {chunk.source}, Page: {chunk.page}, "
                f"Score: {chunk.similarity_score:.3f}\n"
                f"ID: {chunk.chroma_id}\n"
                f"Content: {chunk.content[:600]}\n---"
            )

        prompt = _ROOT_CAUSE_PROMPT.format(
            section=req.section,
            original_section_content=self._extract_section(req.original_report, req.section),
            correction_text=req.correction_text,
            retrieved_chunks=chunks_text,
            features_json=json.dumps(req.features, indent=2),
        )

        raw = self.llm.generate(
            system_prompt="You are a RAG quality analyst. Respond only with valid JSON.",
            user_prompt=prompt,
            temperature=0.1,
            max_tokens=1024,
        )

        match = re.search(r"\{.*\}", raw, re.DOTALL)
        if not match:
            return "UNKNOWN", {}

        try:
            analysis = json.loads(match.group())
        except json.JSONDecodeError:
            return "UNKNOWN", {}

        root_cause = analysis.get("root_cause", "UNKNOWN")
        if root_cause not in ("CHUNK_ERROR", "RETRIEVAL_MISS", "LLM_HALLUCINATION"):
            root_cause = "UNKNOWN"

        return root_cause, analysis

    # -----------------------------------------------------------------------
    # Remediation
    # -----------------------------------------------------------------------

    def _remediate(
        self,
        req: CorrectionRequest,
        root_cause: RootCause,
        analysis: Dict,
        correction_id: str,
    ) -> Tuple[str, List[str]]:
        if root_cause == "CHUNK_ERROR":
            return self._fix_chunk_error(req, analysis, correction_id)
        elif root_cause == "RETRIEVAL_MISS":
            return self._fix_retrieval_miss(req, analysis, correction_id)
        elif root_cause == "LLM_HALLUCINATION":
            return (
                "LLM hallucination identified. No ChromaDB changes made. "
                "Correction logged for prompt engineering review.",
                [],
            )
        else:
            return "Root cause unknown — no ChromaDB changes made.", []

    def _fix_chunk_error(
        self, req: CorrectionRequest, analysis: Dict, correction_id: str
    ) -> Tuple[str, List[str]]:
        problematic_indices: List[int] = analysis.get("problematic_chunk_indices") or [0]

        affected_ids = []
        actions = []

        for idx in problematic_indices:
            if idx >= len(req.retrieved_context):
                continue
            chunk_ref = req.retrieved_context[idx]

            existing = self.kb.get_chunk_by_id(chunk_ref.chroma_id)
            if not existing:
                actions.append(f"Chunk {chunk_ref.chroma_id} not found in DB.")
                continue

            fix_prompt = _CHUNK_FIX_PROMPT.format(
                original_content=existing["content"],
                correction_text=req.correction_text,
                corrected_snippet=req.corrected_snippet or "(none provided)",
            )
            corrected_content = self.llm.generate(
                system_prompt="You are a medical knowledge base editor. Output only the corrected text.",
                user_prompt=fix_prompt,
                temperature=0.1,
                max_tokens=2048,
            )

            self.kb.update_chunk(chunk_ref.chroma_id, corrected_content)
            self.store.save_chunk_edit(
                correction_id=correction_id,
                chroma_id=chunk_ref.chroma_id,
                edit_type="UPDATE",
                original=existing["content"],
                new_content=corrected_content,
            )
            affected_ids.append(chunk_ref.chroma_id)
            actions.append(
                f"Updated chunk {chunk_ref.chroma_id} "
                f"(source: {chunk_ref.source}, page: {chunk_ref.page})"
            )

        return "; ".join(actions) or "No chunks updated.", affected_ids

    def _fix_retrieval_miss(
        self, req: CorrectionRequest, analysis: Dict, correction_id: str
    ) -> Tuple[str, List[str]]:
        new_chunk_text = req.corrected_snippet or analysis.get("new_chunk_suggestion")

        if not new_chunk_text or len(new_chunk_text.strip()) < 50:
            return (
                "RETRIEVAL_MISS identified — insufficient snippet to add new chunk. "
                "Manual review required.",
                [],
            )

        hints = analysis.get("new_chunk_metadata_hints") or {}
        new_metadata = {
            "chunk_id": -1,
            "source": "doctor_correction",
            "page": 0,
            "guideline_source": hints.get("guideline_source", "Unknown"),
            "content_type": hints.get("content_type", "recommendation"),
            "quality_score": 0.9,
            "stone_size_range": hints.get("stone_size_range", ""),
            "stone_location": hints.get("stone_location", ""),
            "evidence_level": hints.get("evidence_level", ""),
            "topics": hints.get("topics", ""),
            "section": "",
            "section_title": f"Doctor correction — {req.section}",
            "treatment_modality": "",
        }

        new_id = self.kb.add_chunk(new_chunk_text, new_metadata)
        self.store.save_chunk_edit(
            correction_id=correction_id,
            chroma_id=new_id,
            edit_type="ADD",
            original=None,
            new_content=new_chunk_text,
        )

        return (
            f"Added new doctor-verified chunk (ID: {new_id}) to address retrieval gap.",
            [new_id],
        )

    # -----------------------------------------------------------------------
    # Utilities
    # -----------------------------------------------------------------------

    def _extract_section(self, report: str, section: str) -> str:
        section_map = {
            "FINDINGS": "1. FINDINGS",
            "IMAGE_CHARACTERISTICS": "2. IMAGE CHARACTERISTICS",
            "CLINICAL_INTERPRETATION": "3. CLINICAL INTERPRETATION",
            "MANAGEMENT_CONSIDERATIONS": "4. MANAGEMENT CONSIDERATIONS",
            "LIMITATIONS": "5. LIMITATIONS",
            "REFERENCES": "6. REFERENCES",
            "GENERAL": None,
        }
        heading = section_map.get(section)
        if not heading:
            return report[:2000]

        pattern = rf"\*\*{re.escape(heading)}\*\*.*?(?=\*\*\d+\.|$)"
        match = re.search(pattern, report, re.DOTALL)
        return match.group(0)[:1500] if match else report[:1500]

    def _build_verification_query(self, req: CorrectionRequest) -> str:
        features = req.features
        size = ""
        stones = features.get("stones") or []
        if stones:
            s = stones[0]
            size_info = s.get("size") or {}
            mm = size_info.get("max_dimension_mm", "")
            if mm:
                size = f"{mm}mm"
        section_label = req.section.lower().replace("_", " ")
        return f"kidney stone {size} {section_label} clinical guideline recommendation".strip()

    def _fix_pending_edit_ids(self, correction_id: str) -> None:
        """chunk_edits tablosundaki PENDING correction_id'leri gerçek ID ile güncelle."""
        import sqlite3
        with sqlite3.connect(self.store.db_path) as conn:
            conn.execute(
                "UPDATE chunk_edits SET correction_id=? WHERE correction_id='PENDING'",
                (correction_id,),
            )

    # -----------------------------------------------------------------------
    # Startup replay
    # -----------------------------------------------------------------------

    def replay_corrections_on_startup(self) -> Dict:
        """
        Uygulama başlarken chunk_edits tablosunu okuyup ChromaDB'ye yeniden uygular.
        HF Spaces gibi ephemeral ortamlarda restart sonrası kayıpları giderir.
        Idempotent: aynı chunk'ı iki kez uygulamaz.
        """
        import sqlite3

        stats = {"replayed_updates": 0, "replayed_adds": 0, "skipped": 0, "errors": 0}

        try:
            with sqlite3.connect(self.store.db_path) as conn:
                conn.row_factory = sqlite3.Row
                edits = conn.execute(
                    "SELECT chroma_id, edit_type, original_content, new_content "
                    "FROM chunk_edits ORDER BY timestamp ASC"
                ).fetchall()
        except Exception as e:
            print(f"[replay] correction_store okunamadı: {e}")
            return stats

        for edit in edits:
            chroma_id = edit["chroma_id"]
            edit_type = edit["edit_type"]
            new_content = edit["new_content"]
            original_content = edit["original_content"]

            try:
                if edit_type == "UPDATE":
                    existing = self.kb.get_chunk_by_id(chroma_id)
                    if existing is None:
                        # Chunk silinmiş olabilir, atla
                        stats["skipped"] += 1
                        continue
                    if existing["content"] == new_content:
                        # Zaten güncel, yeniden uygulamaya gerek yok
                        stats["skipped"] += 1
                        continue
                    self.kb.update_chunk(chroma_id, new_content)
                    stats["replayed_updates"] += 1

                elif edit_type == "ADD":
                    existing = self.kb.get_chunk_by_id(chroma_id)
                    if existing is not None:
                        # Chunk zaten var, tekrar ekleme
                        stats["skipped"] += 1
                        continue
                    # Chunk kaybolmuş, aynı ID ile yeniden ekle
                    embedding = self.kb.embeddings.embed_query(new_content)
                    self.kb.vectorstore._collection.add(
                        ids=[chroma_id],
                        documents=[new_content],
                        embeddings=[embedding],
                        metadatas=[{"source": "doctor_correction", "quality_score": 0.9}],
                    )
                    stats["replayed_adds"] += 1

            except Exception as e:
                print(f"[replay] {edit_type} {chroma_id} uygulanamadı: {e}")
                stats["errors"] += 1

        total = stats["replayed_updates"] + stats["replayed_adds"]
        print(f"[replay] Tamamlandı — {total} düzeltme uygulandı, "
              f"{stats['skipped']} atlandı, {stats['errors']} hata")
        return stats