# -*- coding: utf-8 -*-
"""STEP 6 — Claude client + ai.task ledger (single-shot).

This is a TIER 3 capability. The STRUCTURE is built here; the legal CONTENT it
produces is born as an unapproved draft and is NOT authoritative until a licensed
attorney validates and approves it (Gate 1, Step 3). Nothing here decides law.

Two pieces:
  * familylaw.ai.client — an AbstractModel service wrapping the outbound HTTPS call
    to api.anthropic.com/v1/messages. The provider is pluggable via _route_model()
    (Ollama is a future config flip, not wired). Trust/billing reconciliation NEVER
    reaches a model — the router raises first (IOTA, Bar Rule 5-1.1: human-only).
  * familylaw.ai.task — the audit ledger: model used, token counts, estimated cost,
    and latency are recorded on EVERY call so the economics are measurable and capped.

External calls are ALWAYS mocked in tests (patch _call_provider). The real network
call lives only in _call_provider.
"""

import json
import logging
import time

from odoo import api, fields, models, _
from odoo.exceptions import UserError

_logger = logging.getLogger(__name__)

# ir.config_parameter keys
PARAM_API_KEY = "familylaw.anthropic_api_key"
PARAM_MODEL = "familylaw.model"
PARAM_TOKEN_CEILING = "familylaw.token_ceiling"
PARAM_ENDPOINT = "familylaw.anthropic_endpoint"

DEFAULT_MODEL = "claude-sonnet-4-6"
DEFAULT_ENDPOINT = "https://api.anthropic.com/v1/messages"
DEFAULT_TOKEN_CEILING = 4096
ANTHROPIC_VERSION = "2023-06-01"

# Task types that must NEVER reach any model (human-only, Bar Rule 5-1.1 / IOTA).
BLOCKED_TASK_TYPES = {"trust_billing", "trust_reconciliation", "iota"}

# Estimated prices, USD per 1M tokens. ESTIMATE ONLY — verify current pricing.
MODEL_PRICES = {
    "claude-opus-4-8": {"in": 5.0, "out": 25.0},
    "claude-sonnet-4-6": {"in": 3.0, "out": 15.0},
    "claude-haiku-4-5": {"in": 1.0, "out": 5.0},
}
_DEFAULT_PRICE = {"in": 3.0, "out": 15.0}


class FamilyLawAITask(models.Model):
    _name = "familylaw.ai.task"
    _description = "AI Task Ledger"
    _inherit = ["mail.thread", "mail.activity.mixin"]
    _order = "create_date desc"

    name = fields.Char(required=True, default="AI Task")
    task_type = fields.Char(string="Task Type", required=True, tracking=True)
    provider = fields.Char(string="Provider", tracking=True)
    model_used = fields.Char(string="Model", tracking=True,
                             help="The exact model id used, recorded for audit.")
    state = fields.Selection(
        selection=[
            ("queued", "Queued"),
            ("running", "Running"),
            ("done", "Done"),
            ("failed", "Failed"),
            ("blocked", "Blocked"),
        ],
        default="queued",
        required=True,
        tracking=True,
    )
    case_id = fields.Many2one("familylaw.case", index=True)
    proceeding_id = fields.Many2one("familylaw.proceeding", index=True)
    document_id = fields.Many2one("familylaw.document", index=True)

    prompt_tokens = fields.Integer(string="Prompt Tokens")
    completion_tokens = fields.Integer(string="Completion Tokens")
    total_tokens = fields.Integer(string="Total Tokens")
    token_ceiling = fields.Integer(string="Token Ceiling (max_tokens)")
    cost_usd = fields.Float(string="Est. Cost (USD)", digits=(12, 6))
    latency_ms = fields.Integer(string="Latency (ms)")

    request_summary = fields.Text(string="Request Summary")
    response_summary = fields.Text(string="Response Summary")
    error_message = fields.Text(string="Error")


class FamilyLawAIClient(models.AbstractModel):
    _name = "familylaw.ai.client"
    _description = "Claude API client (single-shot)"

    # --- config -------------------------------------------------------------
    def _icp(self):
        return self.env["ir.config_parameter"].sudo()

    def _get_api_key(self):
        key = self._icp().get_param(PARAM_API_KEY)
        if not key:
            raise UserError(
                _("No Claude API key configured. Set the '%s' system parameter "
                  "(attorney-group readable only).") % PARAM_API_KEY
            )
        return key

    def _get_model(self):
        return self._icp().get_param(PARAM_MODEL) or DEFAULT_MODEL

    def _get_endpoint(self):
        return self._icp().get_param(PARAM_ENDPOINT) or DEFAULT_ENDPOINT

    def _get_token_ceiling(self):
        raw = self._icp().get_param(PARAM_TOKEN_CEILING)
        try:
            return int(raw) if raw else DEFAULT_TOKEN_CEILING
        except (TypeError, ValueError):
            return DEFAULT_TOKEN_CEILING

    # --- routing (pluggable) ------------------------------------------------
    def _route_model(self, task_type, confidential=False):
        """Decide (provider, model). Raises for human-only task types BEFORE any
        network call. Ollama routing for high-confidentiality work is a future
        config flip — the seam is here, not wired."""
        if task_type in BLOCKED_TASK_TYPES:
            raise UserError(
                _("Task type '%s' involves trust/billing reconciliation and must "
                  "never be sent to any AI model (Bar Rule 5-1.1, IOTA — human-only).")
                % task_type
            )
        # Future: if confidential and an Ollama model is configured, route local.
        return "anthropic", self._get_model()

    # --- cost ---------------------------------------------------------------
    def _estimate_cost(self, model, prompt_tokens, completion_tokens):
        price = MODEL_PRICES.get(model, _DEFAULT_PRICE)
        return round(
            (prompt_tokens / 1_000_000.0) * price["in"]
            + (completion_tokens / 1_000_000.0) * price["out"],
            6,
        )

    # --- the actual network call (MOCK THIS IN TESTS) -----------------------
    def _call_provider(self, provider, model, messages, max_tokens, system=None):
        """Perform the outbound HTTPS call. Returns a normalised dict:
            {"text": str, "usage": {"input_tokens": int, "output_tokens": int},
             "citations": [ ... optional ... ]}
        Tests patch this method — they must never hit the network."""
        import requests  # Odoo bundles requests
        if provider != "anthropic":
            raise UserError(_("Provider '%s' is not wired yet.") % provider)
        payload = {
            "model": model,
            "max_tokens": max_tokens,
            "messages": messages,
        }
        if system:
            payload["system"] = system
        headers = {
            "x-api-key": self._get_api_key(),
            "anthropic-version": ANTHROPIC_VERSION,
            "content-type": "application/json",
        }
        resp = requests.post(
            self._get_endpoint(), headers=headers,
            data=json.dumps(payload), timeout=120,
        )
        resp.raise_for_status()
        data = resp.json()
        text = "".join(
            block.get("text", "")
            for block in data.get("content", [])
            if block.get("type") == "text"
        )
        return {"text": text, "usage": data.get("usage", {}), "citations": []}

    # --- high-level entry point ---------------------------------------------
    def generate(self, task_type, messages, *, system=None, case=None,
                 proceeding=None, confidential=False, max_tokens=None,
                 name=None):
        """Route, call, and log an ai.task. Returns (result_dict, task)."""
        provider, model = self._route_model(task_type, confidential)  # may raise
        ceiling = self._get_token_ceiling()
        max_tokens = min(max_tokens or ceiling, ceiling)

        Task = self.env["familylaw.ai.task"]
        task = Task.create({
            "name": name or _("AI: %s") % task_type,
            "task_type": task_type,
            "provider": provider,
            "model_used": model,
            "token_ceiling": max_tokens,
            "state": "running",
            "case_id": case.id if case else False,
            "proceeding_id": proceeding.id if proceeding else False,
            "request_summary": (messages[-1]["content"][:500]
                                if messages else ""),
        })

        start = time.time()
        try:
            result = self._call_provider(provider, model, messages, max_tokens, system)
        except Exception as exc:  # noqa: BLE001 — durably log + re-raise
            # The in-transaction `task` (state 'running') rolls back with the failing
            # transaction, so we record the failure in an INDEPENDENT transaction
            # (durable). The error still propagates — nothing downstream proceeds on
            # a failed/empty result.
            self._log_failure_durably({
                "name": task.name,
                "task_type": task_type,
                "provider": provider,
                "model_used": model,
                "token_ceiling": max_tokens,
                "state": "failed",
                "case_id": case.id if case else False,
                "proceeding_id": proceeding.id if proceeding else False,
                "error_message": str(exc),
                "latency_ms": int((time.time() - start) * 1000),
                "request_summary": (messages[-1]["content"][:500] if messages else ""),
            })
            _logger.warning("AI task '%s' failed (durably logged): %s",
                            task_type, exc)
            raise

        latency = int((time.time() - start) * 1000)
        usage = result.get("usage", {}) or {}
        pt = usage.get("input_tokens", 0) or 0
        ct = usage.get("output_tokens", 0) or 0
        task.write({
            "state": "done",
            "prompt_tokens": pt,
            "completion_tokens": ct,
            "total_tokens": pt + ct,
            "cost_usd": self._estimate_cost(model, pt, ct),
            "latency_ms": latency,
            "response_summary": (result.get("text") or "")[:1000],
        })
        return result, task

    def _log_failure_durably(self, vals):
        """Record a failed ai.task in an INDEPENDENT transaction so the failure
        survives the rollback of the calling transaction (audit/cost ledger).

        In test mode the registry cursor shares the test transaction and commit is a
        savepoint, so the row is visible to the test and cleaned up with it. Audit
        logging must never mask the real error, so all exceptions here are swallowed."""
        try:
            with self.env.registry.cursor() as new_cr:
                new_env = api.Environment(new_cr, self.env.uid, self.env.context)
                new_env["familylaw.ai.task"].create(vals)
                new_cr.commit()
        except Exception:  # noqa: BLE001 — never mask the original failure
            _logger.exception("Could not durably log AI task failure")

    def _purge_task_ids(self, task_ids):
        """Delete ai.task rows in an independent transaction (e.g. pruning old
        failure logs, or test cleanup of durably-committed rows)."""
        if not task_ids:
            return
        try:
            with self.env.registry.cursor() as new_cr:
                new_env = api.Environment(new_cr, self.env.uid, self.env.context)
                new_env["familylaw.ai.task"].browse(task_ids).unlink()
                new_cr.commit()
        except Exception:  # noqa: BLE001
            _logger.exception("Could not purge ai.task rows %s", task_ids)
