"""
CA Automator — CA Assistant chat backend.

Wires the "For CA Agent" architecture (see docs/ca-agent-original/ for the
unmodified source and full design rationale) into this application:

  - system prompt = system_prompt.md (always) plus only the domain reasoning
    module(s) the question actually touches (audit_copilot.md,
    tax_advisory.md, financial_reporting.md, document_templates.md,
    transfer_pricing.md) — this is what system_prompt.md's own "Operating
    Modes" section already documents as the design (each mode names the one
    module it draws on); concatenating every module on every request was
    only ever a v1 shortcut that assumed a large-context paid model. Loading
    just the relevant module(s) also matters practically now that a
    free-tier provider (Groq, capped at a few thousand tokens per minute per
    model) is the default — the full concatenation is ~14,600 tokens, well
    over that budget. Cached per unique domain combination.
  - <knowledge_context> is retrieved live from this firm's own rate_master
    table rather than a second, separate knowledge store — the schema the
    original design proposes (topic/domain/source/effective_date/
    verification_status) maps onto rate_master's own columns almost
    field-for-field, and rate_master is already what every calculation
    module in this app treats as the single source of truth for anything
    statutory.
  - <client_context> is assembled from the selected client + engagement's
    own facts, never carried across a different client's session.
  - Read-only for v1: the assistant can draft a proposed Rate Master entry
    in its reply text when asked (UPDATE KNOWLEDGE etc.), but it does not
    write one itself — the system prompt tells it to hand the fields back
    to the user to enter via the Rate Master screen, which always inserts
    as UNVERIFIED pending human review anyway.

Requires an AI provider configured (GROQ_API_KEY — free, console.groq.com,
no card required — or ANTHROPIC_API_KEY) — see ai_provider.py, the same
dispatcher site_chat.py (the public landing-page chatbot) uses. If neither
is set, chat() raises AssistantError with a plain-language explanation
rather than letting a provider SDK exception surface as a 500.
"""
import os
from datetime import date

import ai_provider
import db

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
PROMPT_DIR = os.path.join(BASE_DIR, "ca_assistant_prompts")

_MAX_TOKENS = 2000

_system_cache = {}

# Which prompt module each detected domain draws its reasoning scaffolding
# from — see system_prompt.md §6 "Operating Modes" for the same mapping,
# stated there for a human/model reader. Several _DOMAIN_KEYWORDS entries
# below share tax_advisory.md, since that module covers direct, indirect
# (VAT) and cross-border tax as one "Tax Advisory" mode.
_DOMAIN_MODULE_FILES = {
    "income_tax": "tax_advisory.md", "tds": "tax_advisory.md", "vat": "tax_advisory.md",
    "depreciation": "tax_advisory.md", "penalty": "tax_advisory.md", "payroll": "tax_advisory.md",
    "grant": "tax_advisory.md", "fdi": "tax_advisory.md",
    "transfer_pricing": "transfer_pricing.md",
    "audit": "audit_copilot.md",
    "financial_reporting": "financial_reporting.md",
    "document_drafting": "document_templates.md",
}


class AssistantError(Exception):
    pass


def _load_system_prompt(domains=None):
    key = tuple(sorted(set(domains or [])))
    if key in _system_cache:
        return _system_cache[key]
    files = ["system_prompt.md"] + sorted(set(
        _DOMAIN_MODULE_FILES[d] for d in key if d in _DOMAIN_MODULE_FILES))
    parts = []
    for fname in files:
        with open(os.path.join(PROMPT_DIR, fname), encoding="utf-8") as f:
            parts.append(f.read())
    _system_cache[key] = "\n\n---\n\n".join(parts)
    return _system_cache[key]


def status():
    return ai_provider.status()


# --------------------------------------------------------- knowledge context
# rate_master domain keywords worth pattern-matching against the user's
# message. Not a real search index — a firm's rate_master is a few hundred
# rows at most, so a keyword scan per request is transparent and cheap, and
# it's easy for a human to audit exactly why an entry was or wasn't pulled in.

_DOMAIN_KEYWORDS = {
    "income_tax": ["income tax", "slab", "instalment", "estimated return", "d-01", "d-03", "business tax", "presumptive"],
    "tds": ["tds", "withholding", "deduct at source"],
    "vat": ["vat", "value added tax", "input tax", "output tax", "reverse charge"],
    "depreciation": ["depreciation", "schedule 2", "pool", "wdv", "written down value"],
    "penalty": ["penalty", "late fee", "interest on", "fine"],
    "payroll": ["salary", "payroll", "ssf", "social security", "gratuity", "provident fund", "rent tds"],
    "grant": ["grant", "subsidy"],
    "fdi": ["fdi", "foreign investment", "foreign direct", "repatriation", "doi"],
    "transfer_pricing": ["transfer pricing", "arm's length", "arms length", "related party", "associated enterprise",
                         "far analysis", "benchmarking", "comparable", "tnmm", "cup method", "resale price method",
                         "cost plus method", "profit split", "form 3ceb", "master file", "local file",
                         "intercompany", "intra-group"],
    "audit": ["audit", "auditor", "materiality", "going concern", "internal control", "substantive procedure",
             "test of controls", "engagement letter", "walkthrough", "sampling", "fraud risk", "vouching",
             "confirmation", "physical verification", "audit opinion", "qualified opinion"],
    "financial_reporting": ["nfrs", "nas ", "ias ", "ifrs", "financial statement", "balance sheet", "recognition",
                            "measurement", "disclosure", "accounting standard", "impairment", "consolidat",
                            "right-of-use", "revenue recognition", "deferred tax", "provision"],
    "document_drafting": ["working paper", "draft a", "draft the", "draft letter", "draft memo", "checklist",
                          "mrl", "management representation", "issue sheet", "engagement letter template"],
}


def _matched_domains(text):
    text = (text or "").lower()
    return [d for d, kws in _DOMAIN_KEYWORDS.items() if any(k in text for k in kws)]


def _current_fy():
    row = db.query("SELECT label_bs FROM fiscal_year WHERE is_current = 1 LIMIT 1", one=True)
    return row["label_bs"] if row else None


def _tp_client_country(client_id):
    if not client_id:
        return None
    row = db.query("SELECT country FROM tp_entity WHERE client_id = ? AND is_the_client = 1 LIMIT 1",
                    (client_id,), one=True)
    return row["country"] if row else None


def build_knowledge_context(user_message, fy=None, max_entries=20, client_id=None):
    domains = _matched_domains(user_message)
    if not domains:
        return ""
    fy = fy or _current_fy()
    rows = []
    for d in domains:
        if d == "transfer_pricing":
            continue  # no rate_master rows in this domain — handled below via tp_regulatory_rule
        rows.extend(db.rates(domain=d, fy=fy))
    rows = rows[:max_entries]

    tp_lines = []
    if "transfer_pricing" in domains:
        country = _tp_client_country(client_id)
        if country:
            tp_rows = db.query(
                "SELECT * FROM tp_regulatory_rule WHERE country = ? AND (financial_year_label = ? OR financial_year_label IS NULL)"
                " ORDER BY domain, code LIMIT ?", (country, fy, max_entries))
            for r in tp_rows:
                value = r["value_num"] if r["value_num"] is not None else (r["value_text"] or "not recorded — see notes")
                unit_suffix = f" {r['unit']}" if r["unit"] and r["unit"] != "text" and r["value_num"] is not None else ""
                stale = " STALE/UNVERIFIED — verify before relying." if r["verification_status"] == "UNVERIFIED" else ""
                tp_lines.append(
                    f"[Topic: {r['label']}] [Jurisdiction: {r['country']}] [Domain: {r['domain']}] [Applies to: {r['applies_to'] or 'all'}]\n"
                    f"Position: {value}{unit_suffix}. Financial year: {r['financial_year_label'] or 'standing rule (not year-specific)'}.\n"
                    f"Law: {r['law_name'] or 'not recorded'} {r['section_article'] or ''}. Authority: {r['source_authority'] or 'not recorded'}.\n"
                    f"Interpretation: {r['interpretation'] or 'not recorded'}.\n"
                    f"Last verified: {r['last_verified_on'] or 'never'} ({r['verification_status']}).{stale}")
        else:
            tp_lines.append(
                "[Transfer pricing jurisdiction not yet selected for this client — no tp_regulatory_rule rows can "
                "be retrieved until the client's country is set in the Transfer Pricing module. Ask for the "
                "jurisdiction before citing any related-party/TP rule.]")

    if not rows and not tp_lines:
        return ""
    lines = ["<knowledge_context>"]
    for r in rows:
        value = r["value_num"] if r["value_num"] is not None else r["value_text"]
        stale = " STALE/UNVERIFIED — verify before relying." if r["verification_status"] == "UNVERIFIED" else ""
        lines.append(
            f"[Topic: {r['label']}] [Domain: {r['domain']}] [Applies to: {r['applies_to'] or 'all'}]\n"
            f"Position: {value} {r['unit'] or ''}. Fiscal year: {r['fiscal_year_label'] or 'not year-specific'}.\n"
            f"Source: {r['source'] or 'not recorded'} {r['clause'] or ''}.\n"
            f"Last verified: {r['last_verified_on'] or 'never'} ({r['verification_status']}).{stale}")
    lines.extend(tp_lines)
    lines.append("</knowledge_context>")
    return "\n".join(lines)


def build_client_context(client_id, engagement_id=None):
    if not client_id:
        return ""
    c = db.query("SELECT * FROM client WHERE id = ?", (client_id,), one=True)
    if not c:
        return ""
    parts = [
        f"Client: {c['name']} ({c['code']})",
        f"Entity type: {c['entity_type'] or 'not recorded'}",
        f"Industry: {c['industry'] or 'not recorded'}",
        "VAT registered: " + ("yes" if c["vat_registered"] else "no")
        + (f", filing frequency: {c['vat_filing_frequency']}" if c.get("vat_filing_frequency") else ""),
        "Special industry: " + ("yes" if c["special_industry"] else "no"),
        f"Stage: {c.get('stage') or 'not recorded'}",
    ]
    if engagement_id:
        e = db.query("SELECT e.*, f.label_bs fy_label FROM engagement e"
                     " LEFT JOIN fiscal_year f ON f.id = e.fiscal_year_id WHERE e.id = ?",
                     (engagement_id,), one=True)
        if e:
            parts.append(f"Engagement: {e['type']} — {e['title'] or ''} — status {e['status']}"
                        + (f" — FY {e['fy_label']}" if e.get("fy_label") else ""))
    return "<client_context>\n" + "\n".join(parts) + "\n</client_context>"


# ------------------------------------------------------------ quick answers
# Client-specific status questions answered straight from this firm's own
# data — no Claude call, no API key, no network dependency at all. These
# are the questions a CA asks every day ("where's this file at", "what's
# still outstanding") and the answer is a database query, not a reasoning
# problem, so there's no reason to spend an API call — or require one to be
# configured — just to read back facts this application already has.

QUICK_QUESTIONS = [
    {"key": "status", "label": "What's the status of this client's file?"},
    {"key": "remaining", "label": "What's remaining / pending?"},
    {"key": "deadlines", "label": "What deadlines are coming up?"},
    {"key": "documents", "label": "What documents are outstanding?"},
    {"key": "billing", "label": "What's the billing status?"},
    {"key": "urgent", "label": "Any urgent items?"},
]


def _days_note(due_date_ad):
    try:
        delta = (date.fromisoformat(due_date_ad) - date.today()).days
    except (TypeError, ValueError):
        return ""
    if delta < 0:
        return f" ({-delta} day{'s' if -delta != 1 else ''} overdue)"
    if delta == 0:
        return " (due today)"
    return f" (due in {delta} day{'s' if delta != 1 else ''})"


def _answer_status(c):
    lines = [f"{c['name']} ({c['code']}) — stage: {c.get('stage') or 'not set'}."]
    assigned = db.query("SELECT full_name FROM app_user WHERE id = ?", (c["assigned_to"],), one=True) \
        if c.get("assigned_to") else None
    lines.append("Looked after by: " + (assigned["full_name"] if assigned else "unassigned") + ".")
    engs = db.query("SELECT e.*, f.label_bs fy_label FROM engagement e"
                    " LEFT JOIN fiscal_year f ON f.id = e.fiscal_year_id"
                    " WHERE e.client_id = ? ORDER BY e.id DESC", (c["id"],))
    if not engs:
        lines.append("No engagements on file yet.")
        return "\n".join(lines)
    lines.append(f"{len(engs)} engagement(s):")
    for e in engs:
        wfs = db.query("SELECT module_key, status FROM workfile WHERE engagement_id = ?", (e["id"],))
        wf_txt = ", ".join(f"{w['module_key']} ({w['status']})" for w in wfs) if wfs else "no workfiles yet"
        lines.append(f"  • {e['type']} — {e['title'] or 'untitled'} — status: {e['status']}"
                    + (f" — FY {e['fy_label']}" if e.get("fy_label") else "")
                    + f". Workfiles: {wf_txt}.")
    return "\n".join(lines)


def _answer_remaining(c):
    followups = db.query("SELECT * FROM client_followup WHERE client_id = ? AND status = 'pending'"
                         " ORDER BY due_date IS NULL, due_date", (c["id"],))
    deadlines = db.query("SELECT * FROM deadline WHERE client_id = ? AND status IN ('pending','late')"
                         " ORDER BY due_date_ad", (c["id"],))
    open_wfs = db.query("SELECT w.module_key, w.status, e.title FROM workfile w"
                        " JOIN engagement e ON e.id = w.engagement_id"
                        " WHERE e.client_id = ? AND w.status != 'signed_off'", (c["id"],))
    if not followups and not deadlines and not open_wfs:
        return f"Nothing outstanding for {c['name']} — no pending follow-ups, reminders, or open workfiles."
    lines = [f"Outstanding for {c['name']}:"]
    if followups:
        lines.append(f"Document follow-ups ({len(followups)}):")
        for f in followups:
            lines.append(f"  • {f['item']}" + (f" — due {f['due_date']}" if f["due_date"] else "")
                        + (f" — {f['note']}" if f["note"] else ""))
    if deadlines:
        lines.append(f"Reminders ({len(deadlines)}):")
        for d in deadlines:
            lines.append(f"  • {d['title']} — due {d['due_date_ad']}{_days_note(d['due_date_ad'])}")
    if open_wfs:
        lines.append(f"Workfiles still open ({len(open_wfs)}):")
        for w in open_wfs:
            lines.append(f"  • {w['module_key']} ({w['status']}) — {w['title'] or ''}")
    return "\n".join(lines)


def _answer_deadlines(c):
    deadlines = db.query("SELECT * FROM deadline WHERE client_id = ? AND status IN ('pending','late')"
                         " ORDER BY due_date_ad", (c["id"],))
    if not deadlines:
        return f"No open reminders for {c['name']}."
    lines = [f"Upcoming/overdue for {c['name']}:"]
    for d in deadlines:
        lines.append(f"  • {d['title']} — {d['due_date_ad']}{_days_note(d['due_date_ad'])}")
    return "\n".join(lines)


def _answer_documents(c):
    followups = db.query("SELECT * FROM client_followup WHERE client_id = ? AND status = 'pending'"
                         " ORDER BY due_date IS NULL, due_date", (c["id"],))
    docs = db.query("SELECT COUNT(*) n FROM document WHERE client_id = ?", (c["id"],), one=True)
    lines = [f"{docs['n']} document(s) on file for {c['name']}."]
    if followups:
        lines.append(f"Still being chased ({len(followups)}):")
        for f in followups:
            lines.append(f"  • {f['item']}" + (f" — due {f['due_date']}" if f["due_date"] else ""))
    else:
        lines.append("Nothing currently being chased from the client.")
    return "\n".join(lines)


def _answer_billing(c):
    invoices = db.query("SELECT * FROM invoice WHERE client_id = ? ORDER BY invoice_date DESC", (c["id"],))
    unbilled = db.query("SELECT t.* FROM time_entry t JOIN engagement e ON e.id = t.engagement_id"
                        " WHERE e.client_id = ? AND t.invoice_id IS NULL AND t.billable = 1", (c["id"],))
    lines = []
    if not invoices:
        lines.append(f"No invoices raised yet for {c['name']}.")
    else:
        outstanding = sum(i["total"] - i["paid_amount"] for i in invoices if i["status"] == "sent")
        lines.append(f"{len(invoices)} invoice(s) for {c['name']}. Outstanding (sent, unpaid): NPR {outstanding:,.2f}.")
        for i in invoices[:5]:
            lines.append(f"  • {i['invoice_no']} — {i['invoice_date']} — {i['status']} — NPR {i['total']:,.2f}")
    if unbilled:
        hours = sum(t["hours"] or 0 for t in unbilled)
        value = sum((t["hours"] or 0) * (t["rate"] or 0) for t in unbilled)
        lines.append(f"Unbilled time: {hours:.1f} hours, approx NPR {value:,.2f} at logged rates.")
    return "\n".join(lines) if lines else f"No billing activity yet for {c['name']}."


def _answer_urgent(c):
    lines = []
    if c.get("urgent"):
        lines.append(f"{c['name']} is flagged URGENT in the Client Directory.")
    today = date.today().isoformat()
    overdue_deadlines = [d for d in db.query(
        "SELECT * FROM deadline WHERE client_id = ? AND status IN ('pending','late')", (c["id"],))
        if d["due_date_ad"] and d["due_date_ad"] < today]
    if overdue_deadlines:
        lines.append(f"{len(overdue_deadlines)} overdue reminder(s):")
        for d in overdue_deadlines:
            lines.append(f"  • {d['title']} — was due {d['due_date_ad']}")
    overdue_invoices = [i for i in db.query(
        "SELECT * FROM invoice WHERE client_id = ? AND status = 'sent'", (c["id"],))
        if i["due_date"] and i["due_date"] < today]
    if overdue_invoices:
        lines.append(f"{len(overdue_invoices)} overdue invoice(s):")
        for i in overdue_invoices:
            lines.append(f"  • {i['invoice_no']} — NPR {i['total']:,.2f} — was due {i['due_date']}")
    if not lines:
        return f"Nothing urgent flagged for {c['name']} right now."
    return "\n".join(lines)


_QUICK_HANDLERS = {
    "status": _answer_status, "remaining": _answer_remaining, "deadlines": _answer_deadlines,
    "documents": _answer_documents, "billing": _answer_billing, "urgent": _answer_urgent,
}


def quick_answer(question_key, client_id):
    if not client_id:
        raise AssistantError("Pick a client first — these questions are client-specific.")
    handler = _QUICK_HANDLERS.get(question_key)
    if not handler:
        raise AssistantError("Unknown question.")
    c = db.query("SELECT * FROM client WHERE id = ?", (client_id,), one=True)
    if not c:
        raise AssistantError("That client no longer exists.")
    return handler(c)


def chat(messages, client_id=None, engagement_id=None, fy=None):
    """messages: [{"role": "user"|"assistant", "content": str}, ...] — the
    full turn history, last entry the new user message (the Messages API is
    stateless; the caller resends history every time). Returns a dict with
    the reply text, the context blocks actually used (for audit logging),
    and token usage."""
    if not ai_provider.active_provider():
        raise AssistantError(
            "No AI provider is configured yet. Set GROQ_API_KEY (free — console.groq.com, no card "
            "required) or ANTHROPIC_API_KEY as an environment variable wherever the server runs, "
            "then restart the server.")
    if not messages or messages[-1].get("role") != "user":
        raise AssistantError("The last message must be from the user")

    last_user = messages[-1]["content"]
    domains = _matched_domains(last_user)
    knowledge_ctx = build_knowledge_context(last_user, fy, client_id=client_id)
    client_ctx = build_client_context(client_id, engagement_id)

    wrapped = list(messages)
    prefix = (client_ctx + "\n" if client_ctx else "") + (knowledge_ctx + "\n" if knowledge_ctx else "")
    if prefix:
        wrapped[-1] = dict(wrapped[-1], content=prefix + "<user_query>\n" + last_user + "\n</user_query>")

    try:
        result = ai_provider.call(
            [{"role": m["role"], "content": m["content"]} for m in wrapped],
            _load_system_prompt(domains), _MAX_TOKENS)
    except ai_provider.ProviderError as e:
        code = str(e)
        messages_by_code = {
            "not_configured": "No AI provider is configured — set GROQ_API_KEY or ANTHROPIC_API_KEY and restart.",
            "invalid_api_key": "The configured AI provider rejected the API key — check it's valid and active.",
            "rate_limited": "Rate limited by the AI provider — try again shortly.",
            "empty_response": "The assistant's reply got cut off mid-thought — try asking again, or more briefly.",
            "unexpected_response": "The AI provider returned an unexpected response — please try again.",
        }
        raise AssistantError(messages_by_code.get(code, f"AI provider error: {code}"))

    return {"text": result["text"], "knowledge_context": knowledge_ctx, "client_context": client_ctx,
            "model": result["model"], "usage": result["usage"]}
