import os
import re
from typing import Any
from azure.identity import DefaultAzureCredential, get_bearer_token_provider
from openai import OpenAI
from pydantic import BaseModel, ConfigDict


class Citation(BaseModel):
    model_config = ConfigDict(extra="forbid")
    chunk_id: str
    quote: str


class Claim(BaseModel):
    model_config = ConfigDict(extra="forbid")
    text: str
    citations: list[Citation]


class GroundedAnswer(BaseModel):
    model_config = ConfigDict(extra="forbid")
    answerable: bool
    claims: list[Claim]


def get_openai_client() -> OpenAI:
    """Initialize OpenAI client configured for Azure OpenAI with Entra ID authentication."""
    endpoint = os.environ.get("AOAI_ENDPOINT", "https://${AZURE_OPENAI_SERVICE}.openai.azure.com")
    token_provider = get_bearer_token_provider(
        DefaultAzureCredential(),
        "https://cognitiveservices.azure.com/.default",
    )
    return OpenAI(
        base_url=f"{endpoint.rstrip('/')}/openai/v1/",
        api_key=token_provider,
    )


SYSTEM_PROMPT = """You are an accurate enterprise assistant.
Answer the user's question using ONLY the provided sources.
For every factual claim you make, include one or more citations specifying:
- chunk_id: the exact chunk ID from which the claim is derived.
- quote: an exact, verbatim text snippet from the chunk supporting the claim.

Rules:
1. Do not use any external knowledge.
2. If the provided sources do not contain enough information to answer the question, set answerable to false and leave claims as an empty list.
3. Every claim must have at least one citation.
4. Quote exact text; do not paraphrase or summarize inside quotes."""


def generate_grounded_answer(
    client: OpenAI,
    model_deployment: str,
    question: str,
    chunks: list[dict[str, Any]],
) -> GroundedAnswer:
    """Generate a structured answer with claim-level citations from source chunks."""
    formatted_sources = "\n\n".join(
        f"[chunk_id: {c['chunk_id']}]\n{c['content']}" for c in chunks
    )
    user_prompt = f"Sources:\n{formatted_sources}\n\nQuestion: {question}"

    completion = client.beta.chat.completions.parse(
        model=model_deployment,
        messages=[
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": user_prompt},
        ],
        response_format=GroundedAnswer,
    )
    return completion.choices[0].message.parsed


def normalize_whitespace(text: str) -> str:
    """Normalise text for quote checks: straight quotes, single spaces, lower case.

    Uses the same rule as verify_citations.py (PB-RAG-06-04), so the live check
    and the incident check give the same result for the same quote.
    """
    text = text.replace("’", "'").replace("‘", "'").replace("“", '"').replace("”", '"')
    return re.sub(r"\s+", " ", text).strip().lower()


def verify_answer_citations(
    answer: GroundedAnswer,
    chunks: list[dict[str, Any]],
) -> tuple[GroundedAnswer, list[str]]:
    """Validate that every cited chunk_id exists and every quote is present in the chunk.

    Returns the verified answer (with invalid citations removed) and a list of verification warnings.
    """
    if not answer.answerable:
        return answer, []

    chunk_map = {c["chunk_id"]: normalize_whitespace(c["content"]) for c in chunks}
    verified_claims: list[Claim] = []
    warnings: list[str] = []

    for claim in answer.claims:
        valid_citations: list[Citation] = []
        for citation in claim.citations:
            if citation.chunk_id not in chunk_map:
                warnings.append(
                    f"Invalid chunk_id '{citation.chunk_id}' in claim: '{claim.text[:40]}...'"
                )
                continue

            chunk_text = chunk_map[citation.chunk_id]
            norm_quote = normalize_whitespace(citation.quote)

            if norm_quote not in chunk_text:
                warnings.append(
                    f"Quote mismatch for chunk '{citation.chunk_id}': '{citation.quote[:30]}...'"
                )
                continue

            valid_citations.append(citation)

        if valid_citations:
            verified_claims.append(Claim(text=claim.text, citations=valid_citations))
        else:
            warnings.append(
                f"Dropped claim with zero valid citations: '{claim.text[:40]}...'"
            )

    is_still_answerable = len(verified_claims) > 0
    return GroundedAnswer(answerable=is_still_answerable, claims=verified_claims), warnings


def build_renderer_contract(
    question: str,
    verified_answer: GroundedAnswer,
    chunks: list[dict[str, Any]],
    raw_answer: GroundedAnswer | None = None,
    warnings: list[str] | None = None,
) -> dict[str, Any]:
    """Map verified claims and chunk metadata to the frontend renderer contract.

    Missing metadata is returned as None, never as a default link or page:
    a made-up link is a renderer fault (PB-RAG-06-04).
    """
    chunk_meta = {c["chunk_id"]: c for c in chunks}
    total_raw_citations = (
        sum(len(cl.citations) for cl in raw_answer.claims) if raw_answer else 0
    )
    verified_citations_count = sum(len(cl.citations) for cl in verified_answer.claims)
    dropped_count = max(0, total_raw_citations - verified_citations_count)

    claims_payload: list[dict[str, Any]] = []
    for claim in verified_answer.claims:
        citations_payload: list[dict[str, Any]] = []
        for cit in claim.citations:
            meta = chunk_meta.get(cit.chunk_id, {})
            citations_payload.append({
                "chunk_id": cit.chunk_id,
                "quote": cit.quote,
                "source_title": meta.get("source_title"),
                "source_url": meta.get("source_url"),
                "page_number": meta.get("page_number"),
            })
        claims_payload.append({
            "text": claim.text,
            "citations": citations_payload,
        })

    status = (
        "passed"
        if not warnings
        else ("passed_with_warnings" if verified_answer.answerable else "rejected")
    )

    return {
        "question": question,
        "answerable": verified_answer.answerable,
        "claims": claims_payload,
        "verification_status": status,
        "citations_verified": verified_citations_count,
        "citations_dropped": dropped_count,
    }


if __name__ == "__main__":
    import json

    print("--- 1. Testing JSON Schema output for structured outputs ---")
    schema = GroundedAnswer.model_json_schema()
    print("Schema title:", schema.get("title"))
    print("Schema required fields:", schema.get("required"))
    print("Citation schema required:", schema["$defs"]["Citation"].get("required"))
    assert schema["$defs"]["Citation"].get("additionalProperties") is False
    assert schema["$defs"]["Claim"].get("additionalProperties") is False
    assert schema.get("additionalProperties") is False

    print("\n--- 2. Testing Citation and Quote Verification logic ---")
    sample_chunks = [
        {
            "chunk_id": "corp-policy-01",
            "content": "All employees must submit business expense claims within 30 calendar days of the transaction date.",
            "source_title": "Corporate Travel and Expense Policy 2026",
            "source_url": "https://${CORP_INTRANET_HOST}/policies/expense-2026.pdf",
            "page_number": 12,
        },
        {
            "chunk_id": "corp-policy-02",
            "content": "Late claims submitted after 60 days require explicit Vice President sign-off and risk rejection.",
            "source_title": "Corporate Travel and Expense Policy 2026",
            "source_url": "https://${CORP_INTRANET_HOST}/policies/expense-2026.pdf",
            "page_number": 13,
        },
    ]

    valid_answer = GroundedAnswer(
        answerable=True,
        claims=[
            Claim(
                text="Expense reports must be submitted within 30 days.",
                citations=[
                    Citation(
                        chunk_id="corp-policy-01",
                        quote="All employees must submit business expense claims within 30 calendar days",
                    )
                ],
            )
        ],
    )
    verified, warn = verify_answer_citations(valid_answer, sample_chunks)
    print("Valid answer verified claims count:", len(verified.claims), "Warnings:", warn)
    assert len(verified.claims) == 1 and len(warn) == 0

    print("\n--- 3. Testing Renderer Contract Mapper ---")
    question_text = "What is the policy for expense report submission deadlines?"
    contract = build_renderer_contract(
        question=question_text,
        verified_answer=verified,
        chunks=sample_chunks,
        raw_answer=valid_answer,
        warnings=warn,
    )
    print("Contract verification_status:", contract["verification_status"])
    print("Contract citations_verified:", contract["citations_verified"])
    assert contract["question"] == question_text
    assert contract["answerable"] is True
    assert contract["verification_status"] == "passed"
    assert contract["citations_verified"] == 1
    assert contract["citations_dropped"] == 0
    assert contract["claims"][0]["citations"][0]["source_title"] == "Corporate Travel and Expense Policy 2026"

    print("\n--- 4. Testing Hallucination / Quote mismatch rejection ---")
    hallucinated_answer = GroundedAnswer(
        answerable=True,
        claims=[
            Claim(
                text="Employees receive double reimbursement on holidays.",
                citations=[
                    Citation(
                        chunk_id="corp-policy-01",
                        quote="Employees receive double reimbursement on holidays",
                    ),
                    Citation(
                        chunk_id="corp-policy-nonexistent",
                        quote="Fabricated text",
                    ),
                ],
            )
        ],
    )
    verified_bad, warn_bad = verify_answer_citations(hallucinated_answer, sample_chunks)
    print("Hallucinated answer verified claims count:", len(verified_bad.claims), "answerable:", verified_bad.answerable)
    print("Captured expected warnings:", warn_bad)
    assert len(verified_bad.claims) == 0
    assert not verified_bad.answerable
    assert len(warn_bad) == 3

    contract_bad = build_renderer_contract(
        question=question_text,
        verified_answer=verified_bad,
        chunks=sample_chunks,
        raw_answer=hallucinated_answer,
        warnings=warn_bad,
    )
    assert contract_bad["answerable"] is False
    assert contract_bad["verification_status"] == "rejected"
    assert contract_bad["citations_dropped"] == 2

    print("\n[SUCCESS] Local self-tests for snippets/answer_with_citations.py passed cleanly!")