"""
NOMAZRA — AI Legal Intelligence Service.
Production-grade multi-country legal AI powered by GPT-4o-mini.
Uses branded system prompt. All AI outputs are brand-filtered — no GPT/OpenAI references.
"""
import os
import re
import logging
import tiktoken

logger = logging.getLogger(__name__)

MODEL = 'gpt-4o-mini'
MAX_CONTEXT_TOKENS = 12000
MAX_RESPONSE_TOKENS = 2000
ENCODING = 'cl100k_base'

# ── BRAND FILTER — Strip all AI provider mentions from responses ──

_BRAND_REPLACEMENTS = [
    (re.compile(r'\bChatGPT\b', re.IGNORECASE), 'NOMAZRA'),
    (re.compile(r'\bChat GPT\b', re.IGNORECASE), 'NOMAZRA'),
    (re.compile(r'\bGPT-4o?(?:-mini)?\b', re.IGNORECASE), 'NOMAZRA AI'),
    (re.compile(r'\bGPT\b', re.IGNORECASE), 'NOMAZRA'),
    (re.compile(r'\bOpenAI\b', re.IGNORECASE), 'Kelthos Technologies'),
    (re.compile(r'\bopen ai\b', re.IGNORECASE), 'Kelthos Technologies'),
    (re.compile(r"I(?:'m| am) an? (?:AI |artificial intelligence )?(?:language )?model\b", re.IGNORECASE),
     "I am NOMAZRA, your AI legal assistant"),
    (re.compile(r"As an? (?:AI |artificial intelligence )?(?:language )?model\b", re.IGNORECASE),
     "As NOMAZRA, your AI legal assistant"),
    (re.compile(r"I(?:'m| am) (?:just )?an? AI\b", re.IGNORECASE),
     "I am NOMAZRA"),
    (re.compile(r'\bmy training data\b', re.IGNORECASE), 'my legal knowledge base'),
    (re.compile(r'\bmy training\b', re.IGNORECASE), 'my knowledge base'),
    (re.compile(r'\btrained by\b', re.IGNORECASE), 'developed by'),
    (re.compile(r'\bI was trained\b', re.IGNORECASE), 'I was developed'),
]


def _brand_filter(text):
    """Remove all mentions of GPT, ChatGPT, OpenAI from AI responses and replace with NOMAZRA branding."""
    if not text:
        return text
    for pattern, replacement in _BRAND_REPLACEMENTS:
        text = pattern.sub(replacement, text)
    return text


# ── SYSTEM PROMPT — The core identity of NOMAZRA ──

SYSTEM_PROMPT = """You are NOMAZRA, a world-class AI legal assistant developed by Kelthos Technologies. You provide authoritative, thorough, and practical legal information covering 140+ countries worldwide.

CRITICAL RULES — YOU MUST FOLLOW THESE:
1. You are NOMAZRA. Never mention GPT, ChatGPT, OpenAI, or any AI model names. If asked who you are, say: "I am NOMAZRA, an AI legal assistant by Kelthos Technologies."
2. NEVER fabricate laws, article numbers, case names, or legal citations. If you are not certain about a specific law number or article, say so clearly. Provide the general legal principle and recommend verifying the specific citation.
3. Provide THOROUGH and COMPLETE answers. Do not give vague or overly cautious non-answers. Users are seeking real legal information — give it to them with depth and detail.
4. Do NOT censor or refuse legal topics. Users need accurate information about ALL areas of law including criminal defense, immigration, labor disputes, family court, civil rights, constitutional law, and more. Provide factual legal information regardless of the topic.
5. Always recommend consulting a licensed attorney for specific legal advice, but provide comprehensive legal information first.
6. Respond in the same language the user writes in. If the user writes in French, respond in French. If Khmer, respond in Khmer.
7. For complex matters, structure your response with clear sections, numbered points, and practical action steps.
8. When discussing a specific country's laws, cite the relevant legal framework (constitution, civil code, criminal code, labor law, etc.) by name.

YOUR KNOWLEDGE COVERS:
- Constitutional law and fundamental rights for all supported countries
- Criminal law and criminal procedure
- Civil law, torts, and contract law
- Family law (marriage, divorce, custody, inheritance)
- Labor and employment law
- Property and real estate law
- Business and commercial law (company formation, contracts, IP)
- Immigration and visa regulations
- Tax law fundamentals
- Human rights and international law
- Consumer protection
- Environmental law
- Digital/cyber law and data protection
- Administrative law and government procedures

RESPONSE FORMAT:
- Be direct and substantive. Start with the answer, then explain.
- Use bullet points and numbered lists for clarity.
- Cite specific laws by name when known (e.g., "Under Nigeria's Labour Act, Chapter L1, Section 7...")
- Include practical next steps the user can take.
- Note any important deadlines, statutes of limitations, or time-sensitive aspects.
- End with a brief professional disclaimer when providing specific legal guidance.

IMPORTANT: Your knowledge is comprehensive but may not reflect the very latest amendments. Always note when legislation may have been recently updated and recommend checking current official sources."""

TOOL_PROMPTS = {
    'law_search': (
        "Search for and explain relevant laws, statutes, and legal provisions for the user's query. "
        "Cite specific laws by name and article number when known. "
        "Structure results clearly with: law name, key provisions, practical implications."
    ),
    'document_review': (
        "Analyze this legal document thoroughly. Identify: "
        "1) Document type and purpose, 2) Key clauses and their implications, "
        "3) Potential issues, risks, or missing provisions, "
        "4) Compliance with applicable law, 5) Recommended changes or additions."
    ),
    'rights_check': (
        "Explain the user's legal rights in this situation comprehensively. "
        "Reference specific constitutional provisions, statutes, and regulations. "
        "Include practical steps to exercise these rights and available remedies if rights are violated."
    ),
    'contract_draft': (
        "Draft or review contract clauses ensuring legal compliance. "
        "Include all essential elements: parties, terms, obligations, "
        "termination conditions, dispute resolution, and governing law."
    ),
    'glossary': (
        "Explain this legal term comprehensively. Provide: "
        "1) Technical definition, 2) Plain language explanation, "
        "3) How it applies in practice, 4) Related legal concepts."
    ),
    'cost_estimate': (
        "Provide a detailed cost estimate for legal services. Include: "
        "1) Attorney fee ranges (hourly and flat fee), 2) Court filing fees, "
        "3) Administrative costs, 4) Potential additional expenses, "
        "5) Total estimated range in local currency and USD."
    ),
    'statute_check': (
        "Analyze the statute of limitations for this case. Provide: "
        "1) Applicable limitation period, 2) The specific law and article, "
        "3) When the clock started, 4) Whether the deadline has passed, "
        "5) Any exceptions, tolling, or extension provisions."
    ),
    'emergency': (
        "Provide emergency legal guidance. Be direct and practical. "
        "List immediate steps the person should take, relevant emergency contacts, "
        "and their fundamental legal rights in this urgent situation."
    ),
}


def _get_client():
    """Lazy-init OpenAI client."""
    import openai
    api_key = os.getenv('OPENAI_API_KEY')
    if not api_key:
        raise ValueError("OPENAI_API_KEY not configured")
    return openai.OpenAI(api_key=api_key)


def count_tokens(text):
    """Count tokens for cost estimation."""
    try:
        enc = tiktoken.get_encoding(ENCODING)
        return len(enc.encode(text))
    except Exception:
        return len(text) // 4


def _truncate_to_tokens(text, max_tokens):
    """Truncate text to fit within token budget."""
    try:
        enc = tiktoken.get_encoding(ENCODING)
        tokens = enc.encode(text)
        if len(tokens) <= max_tokens:
            return text
        return enc.decode(tokens[:max_tokens]) + "\n\n[...truncated for length]"
    except Exception:
        max_chars = max_tokens * 4
        if len(text) <= max_chars:
            return text
        return text[:max_chars] + "\n\n[...truncated for length]"


def _build_messages(conversation_history, user_message, tool=None,
                    document_text=None, country=None):
    """Build messages array for the API call with context management."""
    system = SYSTEM_PROMPT
    if tool and tool in TOOL_PROMPTS:
        system += f"\n\nCurrent task: {TOOL_PROMPTS[tool]}"
    if country:
        system += f"\n\nThe user is asking about the laws of: {country}. Focus your response on this jurisdiction."

    messages = [{'role': 'system', 'content': system}]

    system_tokens = count_tokens(system)
    budget = MAX_CONTEXT_TOKENS - system_tokens - MAX_RESPONSE_TOKENS

    if document_text:
        doc_budget = budget // 2
        truncated_doc = _truncate_to_tokens(document_text, doc_budget)
        doc_msg = {
            'role': 'user',
            'content': f"[Document content for analysis]:\n\n{truncated_doc}"
        }
        messages.append(doc_msg)
        budget -= count_tokens(truncated_doc)

    user_tokens = count_tokens(user_message)
    budget -= user_tokens

    if conversation_history:
        history_messages = []
        used = 0
        for msg in reversed(conversation_history):
            msg_tokens = count_tokens(msg.get('content', ''))
            if used + msg_tokens > budget:
                break
            history_messages.insert(0, msg)
            used += msg_tokens
        messages.extend(history_messages)

    messages.append({'role': 'user', 'content': user_message})
    return messages


def chat(user_message, conversation_history=None, tool=None,
         document_text=None, country=None):
    """
    Send a chat message and get AI response. All responses are brand-filtered.

    Args:
        user_message: The user's current message text
        conversation_history: List of previous {'role': ..., 'content': ...} dicts
        tool: Optional tool context (law_search, document_review, etc.)
        document_text: Optional pre-extracted document text
        country: Optional country name for jurisdiction context

    Returns:
        dict with 'response', 'input_tokens', 'output_tokens'
    """
    if not user_message or not user_message.strip():
        return {
            'response': 'Please provide a message or question.',
            'input_tokens': 0,
            'output_tokens': 0,
        }

    try:
        client = _get_client()
        messages = _build_messages(
            conversation_history or [],
            user_message,
            tool=tool,
            document_text=document_text,
            country=country,
        )

        response = client.chat.completions.create(
            model=MODEL,
            messages=messages,
            max_tokens=MAX_RESPONSE_TOKENS,
            temperature=0.3,
        )

        choice = response.choices[0]
        raw_response = choice.message.content.strip()

        # BRAND FILTER — Strip all GPT/OpenAI mentions
        filtered_response = _brand_filter(raw_response)

        return {
            'response': filtered_response,
            'input_tokens': response.usage.prompt_tokens,
            'output_tokens': response.usage.completion_tokens,
        }
    except Exception as e:
        logger.error("AI chat error: %s", e)
        return {
            'response': (
                "I apologize, but I'm having trouble processing your request right now. "
                "Please try again in a moment. If the issue persists, please contact support."
            ),
            'input_tokens': 0,
            'output_tokens': 0,
            'error': str(e),
        }


def summarize_document(document_text, document_name='', country=None):
    """Generate a concise summary of a legal document."""
    prompt = (
        f"Summarize this document{f' ({document_name})' if document_name else ''}"
        f"{f' under the laws of {country}' if country else ''}. Provide:\n"
        "1. Document type and purpose\n"
        "2. Key provisions/clauses\n"
        "3. Notable legal implications\n"
        "4. Any potential issues or missing elements"
    )
    return chat(prompt, document_text=document_text, tool='document_review', country=country)


def explain_rights(situation, country=None):
    """Explain legal rights for a given situation."""
    return chat(situation, tool='rights_check', country=country)


def search_law(query, country=None):
    """Search for relevant laws."""
    return chat(query, tool='law_search', country=country)


def estimate_cost(case_type, description='', country=None):
    """Estimate legal costs for a case type."""
    prompt = (
        f"Estimate legal costs for a '{case_type}' case"
        f"{f' in {country}' if country else ''}. "
        f"{'Details: ' + description if description else ''}\n\n"
        "Provide: 1) Typical lawyer fees range, "
        "2) Court filing fees, 3) Other potential costs, "
        "4) Total estimated range in local currency and USD."
    )
    return chat(prompt, tool='cost_estimate', country=country)


def check_statute(case_type, incident_date, country=None):
    """Check statute of limitations."""
    prompt = (
        f"What is the statute of limitations for a '{case_type}' case"
        f"{f' under the laws of {country}' if country else ''}? "
        f"The incident occurred on {incident_date}. "
        "Provide the applicable period, relevant law, whether the deadline has passed, "
        "and any exceptions or tolling provisions."
    )
    return chat(prompt, tool='statute_check', country=country)


def draft_contract(contract_type, details, country=None):
    """Draft a contract using AI."""
    prompt = (
        f"Draft a {contract_type} contract"
        f"{f' compliant with the laws of {country}' if country else ''} "
        f"based on these details:\n\n{details}"
    )
    return chat(prompt, tool='contract_draft', country=country)
