core/draft_generator.py

"""
Legal draft generation core.

Two parallel implementations (GPT-4o and Claude) sit behind a single
interface, `DraftGenerator`. A `DraftSession` wraps a generator and
keeps conversation history + refinement count, so stateful UIs only
need a few lines of glue.

Stateless surface (good for Lambda / FastAPI):
    gen = get_generator("claude")           # or "gpt"
    result = gen.generate_initial(...)      # returns InitialDraftResult
    result = gen.refine(message, conversation_history, user_context)

Stateful convenience (good for Streamlit / Gradio):
    session = DraftSession(model="claude")
    result = session.generate_initial(...)
    result = session.refine("add more precedents")
    session.draft_markdown   # latest version
    session.remaining_refinements
"""

from __future__ import annotations

from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import Literal

from .clients import get_anthropic_client, get_openai_client
from .config import get_settings
from .retrieval import format_citations_markdown, search_for_draft
from .templates import parse_json_array


MAX_REFINEMENTS = 5

Message = dict   # {"role": "user" | "assistant", "content": str}


# ===================================================================
# PROMPTS
# ===================================================================

SYSTEM_PROMPT_DRAFT_CLAUDE = """You are an expert legal drafter specialising in Pakistani law.

Hard rules:
- Output ONLY the petition in Markdown. No preamble, no closing remarks, no meta-commentary.
- Never assume facts not given. Use [PLACEHOLDER] where information is missing.
- Follow the structure exactly as instructed in the user turn.
- Integrate every cited law or precedent given to you; do not invent citations.
- Maintain formal, court-ready language throughout."""

SYSTEM_PROMPT_DRAFT_GPT = """You are an expert legal drafter for Pakistani law. Your task is to create a professional, court-ready legal petition in MARKDOWN format using the provided inputs:
1. User Input: Case details including client info, petition type, court, facts, relevant laws, and sections.
2. Knowledge Base Context: Relevant laws, case precedents, and ordinances.
3. Template Draft Analysis: Structure and style guidelines.

Instructions:
- Create a complete, professional legal draft
- Follow the template structure and style
- Integrate legal context with accurate citations
- Use MARKDOWN format with proper headings
- Be comprehensive and court-ready
- ALWAYS adhere to the user's custom instructions/preferences if provided
- Output ONLY the draft, no explanations"""

SYSTEM_PROMPT_REFINE = """You are an expert Pakistani legal drafter continuing to refine a petition.

Hard rules:
- Apply ONLY the changes the user requests; preserve everything else exactly.
- Output the COMPLETE updated petition in Markdown. No explanations before or after.
- Never invent facts or citations.
- Always follow any custom instructions the user has provided."""

SYSTEM_PROMPT_QUERY_CLAUDE = """You are a Pakistani legal research assistant.
Given a case description, produce a JSON array of concise search queries (strings only).
The queries will be run against a vector database containing:
  - The Constitution of Pakistan 1973
  - Punjab High Court and Supreme Court case law
  - FBR Income Tax Ordinance 2001 and Sales Tax Act 1990
Rules:
- Return ONLY a valid JSON array of strings. No markdown fences, no explanation.
- Each query should target a specific legal concept, statute, or precedent relevant to the case.
- Generate between 8 and 15 queries, ordered from most to least specific."""

SYSTEM_PROMPT_QUERY_GPT = (
    "You are a legal research assistant. "
    "A new petition needs to be drafted using the following client/case description. "
    "Devise 5-6 or more concise queries that will be helpful to retrieve relevant information "
    "from a knowledge base containing the Constitution of Pakistan, Punjab case law, "
    "and FBR tax ordinances. "
    "Return ONLY a JSON array of strings, no extra text."
)


# ===================================================================
# RESULT TYPES
# ===================================================================

@dataclass
class InitialDraftResult:
    draft_markdown: str                    # the petition only, no references appended
    draft_with_citations: str              # same draft with `### References` appended (if requested)
    citations: list[dict]                  # [{"score": float, "source": str}, ...]
    context_text: str                      # retrieved KB passages
    template_analysis: str                 # the template style guide that was used
    conversation_history: list[Message]    # carry this into refine()


@dataclass
class RefinementResult:
    draft_markdown: str
    draft_with_citations: str
    conversation_history: list[Message]
    refinement_number: int                 # 1..MAX_REFINEMENTS
    remaining: int                         # MAX_REFINEMENTS - refinement_number


# ===================================================================
# PROMPT BUILDERS (shared)
# ===================================================================

def _build_user_prompt(
    case_text: str,
    context_text: str,
    template_analysis: str,
    user_context: str,
) -> str:
    """Build the rich user-turn that both models consume for the initial draft."""
    custom_section = ""
    if user_context and user_context.strip():
        custom_section = (
            "\n\n<custom_instructions>\n"
            "The user has provided the following persistent preferences. "
            "You MUST follow them throughout this draft and all subsequent refinements:\n"
            f"{user_context.strip()}\n"
            "</custom_instructions>"
        )

    return f"""<task>
Draft a complete, court-ready legal petition in Markdown for a Pakistani court.
</task>
{custom_section}

<case_details>
{case_text}
</case_details>

<knowledge_base_context>
The following passages are retrieved from a vector database of Pakistani law and case precedents.
Integrate the most relevant ones as citations in the appropriate sections of the petition.
{context_text or "(no matches found)"}
</knowledge_base_context>

<template_style_guide>
Follow the structure, tone, length, and formatting described below for the petition.
{template_analysis or "(no template - use standard Pakistani High Court petition format)"}
</template_style_guide>

<required_sections>
1. Title
2. Parties (Petitioner / Respondent)
3. Jurisdiction
4. Statement of Facts
5. Grounds for Petition
6. Legal Arguments (with citations from knowledge base)
7. Prayer / Relief Sought
8. Interim Relief (if applicable)
9. Verification / Affidavit
</required_sections>

Output ONLY the petition in Markdown. No preamble or closing remarks."""


def _build_refinement_prompt(message: str, user_context: str) -> str:
    context_section = ""
    if user_context and user_context.strip():
        context_section = (
            "\n\n**User's persistent custom instructions - always follow these:**\n"
            f"{user_context.strip()}\n"
        )

    return (
        f"The user wants the following change to the petition:{context_section}\n\n"
        f"**Request:** {message}\n\n"
        "Apply the change and return the COMPLETE updated petition in Markdown."
    )


# ===================================================================
# ABSTRACT BASE
# ===================================================================

class DraftGenerator(ABC):
    """Common interface for any LLM-backed draft generator."""

    name: str = "abstract"

    @abstractmethod
    def build_queries(self, case_text: str, max_queries: int = 15) -> list[str]:
        ...

    @abstractmethod
    def _initial_call(self, user_prompt: str) -> tuple[str, list[Message]]:
        """Return (draft_markdown, initial_conversation_history)."""

    @abstractmethod
    def _refine_call(self, conversation_history: list[Message]) -> str:
        """Return updated draft markdown."""

    # ---------- public API ----------

    def generate_initial(
        self,
        case_text: str,
        *,
        user_context: str = "",
        template_analysis: str = "",
        add_citations: bool = True,
        retrieval_top_k: int = 10,
        retrieval_max_chars: int = 10_000,
    ) -> InitialDraftResult:
        if not case_text or not case_text.strip():
            raise ValueError("case_text is required")

        # 1. Knowledge base retrieval
        queries = self.build_queries(case_text)
        context_text, citations = search_for_draft(
            queries,
            top_k=retrieval_top_k,
            max_chars=retrieval_max_chars,
        )

        # 2. Build the user-turn prompt
        user_prompt = _build_user_prompt(
            case_text=case_text,
            context_text=context_text,
            template_analysis=template_analysis,
            user_context=user_context,
        )

        # 3. Call the model
        draft_md, conversation_history = self._initial_call(user_prompt)

        # 4. Optionally append citations for display / docx
        draft_with_citations = draft_md
        if add_citations:
            draft_with_citations += format_citations_markdown(citations)

        return InitialDraftResult(
            draft_markdown=draft_md,
            draft_with_citations=draft_with_citations,
            citations=citations,
            context_text=context_text,
            template_analysis=template_analysis,
            conversation_history=conversation_history,
        )

    def refine(
        self,
        message: str,
        conversation_history: list[Message],
        *,
        user_context: str = "",
        citations: list[dict] | None = None,
        add_citations: bool = True,
        refinement_number: int = 1,
    ) -> RefinementResult:
        if not message or not message.strip():
            raise ValueError("message is required")

        if refinement_number > MAX_REFINEMENTS:
            raise ValueError(
                f"Refinement limit ({MAX_REFINEMENTS}) reached."
            )

        instruction = _build_refinement_prompt(message, user_context)
        new_history = list(conversation_history) + [
            {"role": "user", "content": instruction}
        ]

        draft_md = self._refine_call(new_history)
        new_history.append({"role": "assistant", "content": draft_md})

        draft_with_citations = draft_md
        if add_citations and citations:
            draft_with_citations += format_citations_markdown(citations)

        return RefinementResult(
            draft_markdown=draft_md,
            draft_with_citations=draft_with_citations,
            conversation_history=new_history,
            refinement_number=refinement_number,
            remaining=MAX_REFINEMENTS - refinement_number,
        )


# ===================================================================
# CLAUDE IMPLEMENTATION
# ===================================================================

class ClaudeDraftGenerator(DraftGenerator):
    name = "claude"

    def build_queries(self, case_text: str, max_queries: int = 15) -> list[str]:
        try:
            resp = get_anthropic_client().messages.create(
                model=get_settings().claude_model,
                max_tokens=600,
                system=SYSTEM_PROMPT_QUERY_CLAUDE,
                messages=[{"role": "user", "content": case_text}],
            )
            return parse_json_array(resp.content[0].text.strip(), case_text)[:max_queries]
        except Exception:
            return [case_text[:512]]

    def _initial_call(self, user_prompt: str) -> tuple[str, list[Message]]:
        resp = get_anthropic_client().messages.create(
            model=get_settings().claude_model,
            max_tokens=15_000,
            system=SYSTEM_PROMPT_DRAFT_CLAUDE,
            messages=[{"role": "user", "content": user_prompt}],
        )
        draft_md = resp.content[0].text.strip()
        history = [
            {"role": "user", "content": user_prompt},
            {"role": "assistant", "content": draft_md},
        ]
        return draft_md, history

    def _refine_call(self, conversation_history: list[Message]) -> str:
        resp = get_anthropic_client().messages.create(
            model=get_settings().claude_model,
            max_tokens=15_000,
            system=SYSTEM_PROMPT_REFINE,
            messages=conversation_history,
        )
        return resp.content[0].text.strip()


# ===================================================================
# GPT IMPLEMENTATION
# ===================================================================

class GPTDraftGenerator(DraftGenerator):
    name = "gpt"

    def build_queries(self, case_text: str, max_queries: int = 15) -> list[str]:
        try:
            resp = get_openai_client().chat.completions.create(
                model=get_settings().gpt_model,
                messages=[
                    {"role": "system", "content": SYSTEM_PROMPT_QUERY_GPT},
                    {"role": "user", "content": case_text},
                ],
                temperature=0.1,
                max_tokens=2000,
            )
            return parse_json_array(resp.choices[0].message.content.strip(), case_text)[:max_queries]
        except Exception:
            return [case_text[:512]]

    def _initial_call(self, user_prompt: str) -> tuple[str, list[Message]]:
        # The GPT version keeps the system message in the conversation history
        # so that subsequent refinement turns retain the same instructions.
        history = [
            {"role": "system", "content": SYSTEM_PROMPT_DRAFT_GPT},
            {"role": "user", "content": user_prompt},
        ]
        resp = get_openai_client().chat.completions.create(
            model=get_settings().gpt_draft_model,
            messages=history,
            max_tokens=15_000,
        )
        draft_md = resp.choices[0].message.content.strip()
        history.append({"role": "assistant", "content": draft_md})
        return draft_md, history

    def _refine_call(self, conversation_history: list[Message]) -> str:
        resp = get_openai_client().chat.completions.create(
            model=get_settings().gpt_draft_model,
            messages=conversation_history,
            max_tokens=15_000,
        )
        return resp.choices[0].message.content.strip()


# ===================================================================
# FACTORY
# ===================================================================

ModelName = Literal["gpt", "claude"]


def get_generator(model: ModelName) -> DraftGenerator:
    if model == "claude":
        return ClaudeDraftGenerator()
    if model == "gpt":
        return GPTDraftGenerator()
    raise ValueError(f"Unknown model: {model!r}")


# ===================================================================
# STATEFUL CONVENIENCE WRAPPER
# ===================================================================

@dataclass
class DraftSession:
    """
    Holds session state for a single draft conversation.

    Usage:
        s = DraftSession(model="claude", user_context="Always cite the Constitution.")
        s.generate_initial(case_text="...", template_analysis="...")
        s.refine("Add more precedents")
        s.refine("Shorten the facts section")
        print(s.draft_with_citations)
    """
    model: ModelName = "claude"
    user_context: str = ""
    add_citations: bool = True

    # Populated by generate_initial:
    case_text: str = ""
    citations: list[dict] = field(default_factory=list)
    context_text: str = ""
    template_analysis: str = ""
    draft_markdown: str = ""
    draft_with_citations: str = ""
    conversation_history: list[Message] = field(default_factory=list)
    user_message_count: int = 0

    def __post_init__(self):
        self._generator = get_generator(self.model)

    # ---------- properties ----------

    @property
    def remaining_refinements(self) -> int:
        return MAX_REFINEMENTS - self.user_message_count

    @property
    def at_refinement_limit(self) -> bool:
        return self.user_message_count >= MAX_REFINEMENTS

    # ---------- actions ----------

    def generate_initial(
        self,
        case_text: str,
        *,
        template_analysis: str = "",
    ) -> InitialDraftResult:
        result = self._generator.generate_initial(
            case_text=case_text,
            user_context=self.user_context,
            template_analysis=template_analysis,
            add_citations=self.add_citations,
        )
        self.case_text = case_text
        self.template_analysis = template_analysis
        self.citations = result.citations
        self.context_text = result.context_text
        self.draft_markdown = result.draft_markdown
        self.draft_with_citations = result.draft_with_citations
        self.conversation_history = result.conversation_history
        self.user_message_count = 0
        return result

    def refine(self, message: str) -> RefinementResult:
        if self.at_refinement_limit:
            raise RuntimeError(
                f"Refinement limit ({MAX_REFINEMENTS}) reached. "
                "Start a new session to continue."
            )

        self.user_message_count += 1
        result = self._generator.refine(
            message=message,
            conversation_history=self.conversation_history,
            user_context=self.user_context,
            citations=self.citations,
            add_citations=self.add_citations,
            refinement_number=self.user_message_count,
        )
        self.draft_markdown = result.draft_markdown
        self.draft_with_citations = result.draft_with_citations
        self.conversation_history = result.conversation_history
        return result

    # ---------- serialisation (useful for FastAPI / DynamoDB) ----------

    def to_dict(self) -> dict:
        return {
            "model": self.model,
            "user_context": self.user_context,
            "add_citations": self.add_citations,
            "case_text": self.case_text,
            "citations": self.citations,
            "context_text": self.context_text,
            "template_analysis": self.template_analysis,
            "draft_markdown": self.draft_markdown,
            "draft_with_citations": self.draft_with_citations,
            "conversation_history": self.conversation_history,
            "user_message_count": self.user_message_count,
        }

    @classmethod
    def from_dict(cls, data: dict) -> "DraftSession":
        s = cls(
            model=data["model"],
            user_context=data.get("user_context", ""),
            add_citations=data.get("add_citations", True),
        )
        s.case_text = data.get("case_text", "")
        s.citations = data.get("citations", [])
        s.context_text = data.get("context_text", "")
        s.template_analysis = data.get("template_analysis", "")
        s.draft_markdown = data.get("draft_markdown", "")
        s.draft_with_citations = data.get("draft_with_citations", "")
        s.conversation_history = data.get("conversation_history", [])
        s.user_message_count = data.get("user_message_count", 0)
        return s