"""Season outcome capture for the Maminda KCS learning loop.

Standard library only, so the Gradio app environment can import it.

Predictions are never labels. A model can only learn from two observed outcomes:
  * harvest_outcomes   - measured harvest for one field-season (gives the yield label)
  * repayment_outcomes - what happened to one loan that was scored by a saved prediction
Both are linked to the records the model actually saw, so training rows use the exact
inputs that were available at prediction time (no look-ahead leakage).
"""
from __future__ import annotations
from contextlib import contextmanager
from datetime import datetime, timezone
from pathlib import Path
import json
import math
import sqlite3
import uuid

HARVEST_METHODS = ('weighed_full_field', 'weighed_sample_plots', 'buyer_receipt', 'farmer_estimate')
# Only these methods produce training labels. A farmer's own estimate is stored but not learned from.
LABEL_GRADE_METHODS = ('weighed_full_field', 'weighed_sample_plots', 'buyer_receipt')
FINAL_STATUS = ('paid_on_time', 'paid_late', 'defaulted', 'restructured', 'written_off')

SCHEMA = '''
CREATE TABLE IF NOT EXISTS harvest_outcomes (
  id TEXT PRIMARY KEY,
  context_key TEXT NOT NULL UNIQUE REFERENCES field_seasons(context_key),
  harvest_date TEXT NOT NULL,
  area_ha REAL NOT NULL CHECK(area_ha > 0),
  harvested_kg REAL NOT NULL CHECK(harvested_kg >= 0),
  yield_t_ha REAL NOT NULL CHECK(yield_t_ha >= 0),
  measurement_method TEXT NOT NULL,
  label_grade INTEGER NOT NULL CHECK(label_grade IN (0,1)),
  evidence_reference TEXT NOT NULL,
  recorded_by TEXT NOT NULL,
  recorded_at TEXT NOT NULL);
CREATE TABLE IF NOT EXISTS repayment_outcomes (
  id TEXT PRIMARY KEY,
  prediction_id TEXT NOT NULL UNIQUE REFERENCES research_predictions(id),
  context_key TEXT NOT NULL REFERENCES field_seasons(context_key),
  loan_reference TEXT NOT NULL,
  due_date TEXT NOT NULL,
  amount_due REAL NOT NULL CHECK(amount_due > 0),
  amount_paid_by_due REAL NOT NULL CHECK(amount_paid_by_due >= 0),
  currency TEXT NOT NULL,
  paid_in_full_by_due INTEGER NOT NULL CHECK(paid_in_full_by_due IN (0,1)),
  final_status TEXT NOT NULL,
  evidence_reference TEXT NOT NULL,
  recorded_by TEXT NOT NULL,
  recorded_at TEXT NOT NULL);
CREATE TABLE IF NOT EXISTS outcome_corrections (
  id TEXT PRIMARY KEY, table_name TEXT NOT NULL, record_id TEXT NOT NULL,
  previous_payload TEXT NOT NULL, reason TEXT NOT NULL, corrected_at TEXT NOT NULL);
'''


def utc_now():
    return datetime.now(timezone.utc)


def parse_time(value, name):
    if value is None or str(value).strip() == '':
        raise ValueError(f'Enter the {name}.')
    text = str(value).strip()
    if len(text) == 10:  # plain date: treat as end of that day in Harare (UTC+2)
        text += 'T23:59:59+02:00'
    result = datetime.fromisoformat(text.replace('Z', '+00:00'))
    if result.tzinfo is None:
        raise ValueError(f'Include the time zone in the {name}, e.g. 2027-04-20T12:00:00+02:00, or give a plain date.')
    result = result.astimezone(timezone.utc)
    if result > utc_now():
        raise ValueError(f'The {name} cannot be in the future.')
    return result


def positive(value, name, allow_zero=False):
    if value is None or isinstance(value, bool):
        raise ValueError(f'Enter {name}.')
    value = float(value)
    if not math.isfinite(value) or value < 0 or (value == 0 and not allow_zero):
        raise ValueError(f'{name} must be a finite number above zero.' if not allow_zero else f'{name} cannot be negative.')
    return value


def text(value, name):
    value = '' if value is None else str(value).strip()
    if not value or len(value) > 300:
        raise ValueError(f'Enter {name} (at most 300 characters).')
    return value


class OutcomeStore:
    def __init__(self, database):
        self.path = Path(database)
        with self.connect() as db:
            db.executescript(SCHEMA)

    @contextmanager
    def connect(self):
        db = sqlite3.connect(self.path, timeout=30)
        db.row_factory = sqlite3.Row
        db.execute('PRAGMA foreign_keys=ON')
        try:
            with db:
                yield db
        finally:
            db.close()

    def _audit(self, db, table, row, reason):
        db.execute('INSERT INTO outcome_corrections VALUES (?,?,?,?,?,?)',
                   (str(uuid.uuid4()), table, row['id'], json.dumps(dict(row)), text(reason, 'a correction reason'),
                    utc_now().isoformat()))

    def record_harvest(self, context_key, harvest_date, area_ha, harvested_kg, measurement_method,
                       evidence_reference, recorded_by, correction_reason=None):
        if measurement_method not in HARVEST_METHODS:
            raise ValueError('Measurement method must be one of: ' + ', '.join(HARVEST_METHODS))
        when = parse_time(harvest_date, 'harvest date')
        area = positive(area_ha, 'field area (ha)')
        kg = positive(harvested_kg, 'harvested mass (kg)', allow_zero=True)
        yield_t_ha = kg / 1000 / area
        if yield_t_ha > 150:
            raise ValueError(f'{yield_t_ha:.1f} t/ha is implausible. Check that the mass is in kg and the area in hectares.')
        row = (str(uuid.uuid4()), context_key, when.isoformat(), area, kg, yield_t_ha, measurement_method,
               int(measurement_method in LABEL_GRADE_METHODS), text(evidence_reference, 'an evidence reference (scale sheet, receipt number)'),
               text(recorded_by, 'who recorded it'), utc_now().isoformat())
        with self.connect() as db:
            if not db.execute('SELECT 1 FROM field_seasons WHERE context_key=?', (context_key,)).fetchone():
                raise ValueError('Unknown field-season. Save at least one observation for it first.')
            existing = db.execute('SELECT * FROM harvest_outcomes WHERE context_key=?', (context_key,)).fetchone()
            if existing:
                if not correction_reason:
                    raise ValueError('A harvest is already recorded for this field-season. Give a correction reason to replace it.')
                self._audit(db, 'harvest_outcomes', existing, correction_reason)
                db.execute('DELETE FROM harvest_outcomes WHERE id=?', (existing['id'],))
            db.execute('INSERT INTO harvest_outcomes VALUES (?,?,?,?,?,?,?,?,?,?,?)', row)
        return dict(harvest_id=row[0], yield_t_ha=yield_t_ha, label_grade=bool(row[7]))

    def record_repayment(self, prediction_id, loan_reference, due_date, amount_due, amount_paid_by_due,
                         currency, final_status, evidence_reference, recorded_by, correction_reason=None):
        if final_status not in FINAL_STATUS:
            raise ValueError('Final status must be one of: ' + ', '.join(FINAL_STATUS))
        due = parse_time(due_date, 'due date')
        owed = positive(amount_due, 'amount due')
        paid = positive(amount_paid_by_due, 'amount paid by the due date', allow_zero=True)
        currency = text(currency, 'the currency').upper()
        paid_in_full = int(paid >= owed * 0.999)
        if paid_in_full and final_status in ('defaulted', 'written_off', 'paid_late'):
            raise ValueError('The amount paid by the due date covers the loan, which contradicts the final status.')
        with self.connect() as db:
            prediction = db.execute('SELECT * FROM research_predictions WHERE id=?', (prediction_id,)).fetchone()
            if not prediction:
                raise ValueError('Unknown prediction ID. Use the ID of the KCS estimate used when the loan was issued.')
            payload = json.loads(prediction['payload'])
            credit = payload.get('research_credit_estimate') or {}
            if not credit.get('credit_features'):
                raise ValueError('That prediction has no credit inputs saved, so it cannot become a repayment training example.')
            if datetime.fromisoformat(prediction['created_at']) > due:
                raise ValueError('The KCS estimate was created after the due date. Only pre-loan estimates can be linked.')
            fin = credit.get('financial_inputs', {})
            if fin.get('currency') and fin['currency'] != currency:
                raise ValueError(f"The estimate used {fin['currency']}; record the repayment in the same currency.")
            existing = db.execute('SELECT * FROM repayment_outcomes WHERE prediction_id=?', (prediction_id,)).fetchone()
            if existing:
                if not correction_reason:
                    raise ValueError('An outcome is already recorded for this loan. Give a correction reason to replace it.')
                self._audit(db, 'repayment_outcomes', existing, correction_reason)
                db.execute('DELETE FROM repayment_outcomes WHERE id=?', (existing['id'],))
            row = (str(uuid.uuid4()), prediction_id, prediction['context_key'], text(loan_reference, 'the loan reference'),
                   due.isoformat(), owed, paid, currency, paid_in_full, final_status,
                   text(evidence_reference, 'an evidence reference (statement or ledger ID)'),
                   text(recorded_by, 'who recorded it'), utc_now().isoformat())
            db.execute('INSERT INTO repayment_outcomes VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?)', row)
        return dict(repayment_id=row[0], paid_in_full_by_due=bool(paid_in_full))

    def counts(self):
        with self.connect() as db:
            harvest = db.execute('SELECT count(*), coalesce(sum(label_grade),0) FROM harvest_outcomes').fetchone()
            loans = db.execute('SELECT count(*), coalesce(sum(paid_in_full_by_due),0) FROM repayment_outcomes').fetchone()
        return dict(harvests_recorded=harvest[0], harvests_usable_as_labels=harvest[1],
                    repayments_recorded=loans[0], repaid_in_full_by_due=loans[1])

    def harvests(self, context_key):
        with self.connect() as db:
            return [dict(r) for r in db.execute('SELECT * FROM harvest_outcomes WHERE context_key=?', (context_key,))]