"""
Production daily counter persistence for edge sack zone counter.
Tracks per-zone entries and the daily total per counting day,
delimited by the daily cutoff time. SQLite schema + current_counter.json state.

Zone ids are stable keys like zone_1, zone_2, … (from zones_config).
"""
import json
import sqlite3
import threading
import time
from datetime import datetime, timedelta
from pathlib import Path


class CounterStore:
    def __init__(
        self,
        db_path,
        state_file,
        camera_name,
        object_label='object',
        cutoff_time='20:00',
        carry_ids=50,
        zone_ids=None,
        logger=print,
    ):
        self.db_path = db_path
        self.state_file = Path(state_file)
        self.camera_name = camera_name
        self.object_label = object_label
        self.cutoff_time_str = cutoff_time
        datetime.strptime(cutoff_time, '%H:%M')

        self.carry_ids = int(carry_ids)
        self.zone_ids = list(zone_ids) if zone_ids else ['zone_1']
        if not self.zone_ids:
            raise ValueError('zone_ids must contain at least one zone id')
        self._zone_id_set = set(self.zone_ids)
        self.log = logger

        self.state_lock = threading.Lock()
        self.shutdown_event = threading.Event()

        Path(db_path).parent.mkdir(parents=True, exist_ok=True)
        self.state_file.parent.mkdir(parents=True, exist_ok=True)

        self.db = sqlite3.connect(db_path, check_same_thread=False)
        self._init_db()
        self.current_state = self._load_state()
        with self.state_lock:
            self._fill_missing_days()
            counting_date = self.get_counting_date()
            if self.current_state is None:
                self._start_new_day(counting_date)
            self._persist_day()

    def _table_columns(self, table):
        cur = self.db.cursor()
        cur.execute(f'PRAGMA table_info({table})')
        return {row[1] for row in cur.fetchall()}

    def _init_db(self):
        cur = self.db.cursor()
        cur.execute(
            """
            CREATE TABLE IF NOT EXISTS daily_counters (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                counting_date TEXT NOT NULL,
                camera_name TEXT NOT NULL,
                object_label TEXT NOT NULL,
                total_count INTEGER NOT NULL DEFAULT 0,
                start_time TEXT,
                end_time TEXT,
                updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                UNIQUE(counting_date, camera_name, object_label)
            )
            """
        )
        cur.execute(
            """
            CREATE TABLE IF NOT EXISTS daily_zone_counts (
                counting_date TEXT NOT NULL,
                camera_name TEXT NOT NULL,
                object_label TEXT NOT NULL,
                zone_id TEXT NOT NULL,
                count INTEGER NOT NULL DEFAULT 0,
                PRIMARY KEY (counting_date, camera_name, object_label, zone_id)
            )
            """
        )
        self.db.commit()
        self._migrate_left_right_if_needed()

    def _migrate_left_right_if_needed(self):
        cols = self._table_columns('daily_counters')
        if 'total_left' not in cols and 'total_right' not in cols:
            return

        self.log('Migrating daily_counters left/right columns to daily_zone_counts')
        cur = self.db.cursor()
        cur.execute(
            """
            SELECT counting_date, camera_name, object_label,
                   COALESCE(total_left, 0), COALESCE(total_right, 0)
            FROM daily_counters
            """
        )
        rows = cur.fetchall()
        for counting_date, camera_name, object_label, left, right in rows:
            cur.execute(
                """
                INSERT INTO daily_zone_counts
                    (counting_date, camera_name, object_label, zone_id, count)
                VALUES (?, ?, ?, 'zone_1', ?)
                ON CONFLICT(counting_date, camera_name, object_label, zone_id)
                DO UPDATE SET count = excluded.count
                """,
                (counting_date, camera_name, object_label, left),
            )
            cur.execute(
                """
                INSERT INTO daily_zone_counts
                    (counting_date, camera_name, object_label, zone_id, count)
                VALUES (?, ?, ?, 'zone_2', ?)
                ON CONFLICT(counting_date, camera_name, object_label, zone_id)
                DO UPDATE SET count = excluded.count
                """,
                (counting_date, camera_name, object_label, right),
            )

        cur.execute(
            """
            CREATE TABLE daily_counters_new (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                counting_date TEXT NOT NULL,
                camera_name TEXT NOT NULL,
                object_label TEXT NOT NULL,
                total_count INTEGER NOT NULL DEFAULT 0,
                start_time TEXT,
                end_time TEXT,
                updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                UNIQUE(counting_date, camera_name, object_label)
            )
            """
        )
        cur.execute(
            """
            INSERT INTO daily_counters_new
                (id, counting_date, camera_name, object_label,
                 total_count, start_time, end_time, updated_at)
            SELECT id, counting_date, camera_name, object_label,
                   total_count, start_time, end_time, updated_at
            FROM daily_counters
            """
        )
        cur.execute('DROP TABLE daily_counters')
        cur.execute('ALTER TABLE daily_counters_new RENAME TO daily_counters')
        self.db.commit()
        self.log(f'Migration complete ({len(rows)} daily row(s))')

    def get_counting_date(self, dt=None):
        if dt is None:
            dt = datetime.now()
        cutoff = datetime.strptime(self.cutoff_time_str, '%H:%M').time()
        if dt.time() < cutoff:
            return dt.date().isoformat()
        return (dt.date() + timedelta(days=1)).isoformat()

    def _empty_counts(self):
        return {zid: 0 for zid in self.zone_ids}

    def _normalize_state(self, state):
        counts = state.get('counts')
        if not isinstance(counts, dict):
            counts = {}
            if 'count_left' in state or 'count_right' in state:
                counts['zone_1'] = int(state.get('count_left', 0) or 0)
                counts['zone_2'] = int(state.get('count_right', 0) or 0)
        normalized = self._empty_counts()
        for zid in self.zone_ids:
            normalized[zid] = int(counts.get(zid, 0) or 0)
        # Preserve migrated zone_2 even if current config is 1 zone? Spec says
        # configured ids only — drop unknown keys from active counts.
        state['counts'] = normalized
        state['total_count'] = sum(normalized.values())
        state.pop('count', None)
        state.pop('count_left', None)
        state.pop('count_right', None)
        state.setdefault('counted_event_ids', [])
        state.setdefault('start_time', datetime.now().isoformat())
        state.setdefault('last_detection_time', state.get('start_time'))
        return state

    def _load_state(self):
        if not self.state_file.exists():
            return None
        try:
            with open(self.state_file, 'r', encoding='utf-8') as f:
                state = json.load(f)
            current_date = self.get_counting_date()
            if state.get('counting_date') != current_date:
                self.log(
                    f"State file belongs to previous counting day "
                    f"({state.get('counting_date')}). Starting fresh."
                )
                self.state_file.unlink(missing_ok=True)
                return None
            state = self._normalize_state(state)
            stale_ids = state.get('counted_event_ids') or []
            if stale_ids:
                self.log(
                    f"Cleared {len(stale_ids)} persisted track dedup keys "
                    f"(track IDs reset on restart)"
                )
            state['counted_event_ids'] = []
            self.log(
                f"Resumed {current_date} with total={state['total_count']} "
                f"counts={state['counts']}"
            )
            with open(self.state_file, 'w', encoding='utf-8') as f:
                json.dump(state, f, indent=2, ensure_ascii=False)
            return state
        except Exception as exc:
            self.log(f'Failed to load state file: {exc}')
            return None

    def save_state(self):
        if self.current_state is None:
            self.state_file.unlink(missing_ok=True)
            return
        with open(self.state_file, 'w', encoding='utf-8') as f:
            json.dump(self.current_state, f, indent=2, ensure_ascii=False)

    def _start_new_day(self, counting_date):
        now = datetime.now().isoformat()
        carried = []
        if self.current_state is not None:
            try:
                carried = self.current_state['counted_event_ids'][-self.carry_ids:]
            except (KeyError, TypeError):
                carried = []
        counts = self._empty_counts()
        self.current_state = {
            'counting_date': counting_date,
            'total_count': 0,
            'counts': counts,
            'start_time': now,
            'last_detection_time': now,
            'counted_event_ids': carried,
        }
        self.save_state()
        self.log(f'Started counting day {counting_date} ({self.object_label})')

    def _ensure_day_row(self, counting_date, start_time=None, end_time=None):
        """Insert a zero row for counting_date if it does not already exist."""
        now = datetime.now().isoformat()
        cur = self.db.cursor()
        cur.execute(
            """
            INSERT OR IGNORE INTO daily_counters
                (counting_date, camera_name, object_label,
                 total_count, start_time, end_time)
            VALUES (?, ?, ?, 0, ?, ?)
            """,
            (
                counting_date,
                self.camera_name,
                self.object_label,
                start_time or now,
                end_time or now,
            ),
        )
        for zid in self.zone_ids:
            cur.execute(
                """
                INSERT OR IGNORE INTO daily_zone_counts
                    (counting_date, camera_name, object_label, zone_id, count)
                VALUES (?, ?, ?, ?, 0)
                """,
                (counting_date, self.camera_name, self.object_label, zid),
            )
        self.db.commit()

    def _fill_missing_days(self):
        """Backfill any missing counting dates from first DB row through today as 0."""
        today = self.get_counting_date()
        cur = self.db.cursor()
        cur.execute(
            """
            SELECT counting_date FROM daily_counters
            WHERE camera_name = ? AND object_label = ?
            ORDER BY counting_date ASC
            """,
            (self.camera_name, self.object_label),
        )
        existing = {row[0] for row in cur.fetchall()}
        if not existing:
            self._ensure_day_row(today)
            return

        start = datetime.strptime(min(existing), '%Y-%m-%d').date()
        end = datetime.strptime(today, '%Y-%m-%d').date()
        filled = 0
        d = start
        while d <= end:
            key = d.isoformat()
            if key not in existing:
                self._ensure_day_row(key)
                filled += 1
            d += timedelta(days=1)
        if filled:
            self.log(
                f'Backfilled {filled} zero-activity day(s) '
                f'{start.isoformat()}..{end.isoformat()}'
            )

    def record_zone_entry(self, track_id, zone_id):
        """Record a sack entering a zone. zone_id: 'zone_1' | 'zone_2' | …

        Returns (total_count, day_started, counted).
        """
        if zone_id not in self._zone_id_set:
            raise ValueError(
                f"zone_id must be one of {self.zone_ids}, got {zone_id!r}"
            )
        with self.state_lock:
            counting_date = self.get_counting_date()
            day_started = False
            if self.current_state is None or self.current_state['counting_date'] != counting_date:
                self._start_new_day(counting_date)
                day_started = True

            counted = False
            event_key = f"{track_id}_{zone_id}"
            if event_key not in self.current_state['counted_event_ids']:
                self.current_state['total_count'] += 1
                self.current_state['counts'][zone_id] = (
                    int(self.current_state['counts'].get(zone_id, 0)) + 1
                )
                self.current_state['counted_event_ids'].append(event_key)
                counted = True
                self.log(
                    f'Counted {zone_id} (track {track_id}) | {counting_date} '
                    f'total: {self.current_state["total_count"]} '
                    f'counts={self.current_state["counts"]}'
                )
                self._persist_day()

            self.current_state['last_detection_time'] = datetime.now().isoformat()
            self.save_state()
            return self.current_state['total_count'], day_started, counted

    def _persist_day(self):
        state = self.current_state
        cur = self.db.cursor()
        cur.execute(
            """
            INSERT INTO daily_counters
                (counting_date, camera_name, object_label,
                 total_count, start_time, end_time)
            VALUES (?, ?, ?, ?, ?, ?)
            ON CONFLICT(counting_date, camera_name, object_label)
            DO UPDATE SET
                total_count = excluded.total_count,
                end_time = excluded.end_time,
                updated_at = CURRENT_TIMESTAMP
            """,
            (
                state['counting_date'], self.camera_name, self.object_label,
                state['total_count'],
                state['start_time'], datetime.now().isoformat(),
            ),
        )
        for zid, count in state['counts'].items():
            cur.execute(
                """
                INSERT INTO daily_zone_counts
                    (counting_date, camera_name, object_label, zone_id, count)
                VALUES (?, ?, ?, ?, ?)
                ON CONFLICT(counting_date, camera_name, object_label, zone_id)
                DO UPDATE SET count = excluded.count
                """,
                (
                    state['counting_date'], self.camera_name, self.object_label,
                    zid, int(count),
                ),
            )
        self.db.commit()

    def cutoff_watcher_loop(self):
        while not self.shutdown_event.is_set():
            time.sleep(60)
            with self.state_lock:
                self._fill_missing_days()
                counting_date = self.get_counting_date()
                if self.current_state is None:
                    self._start_new_day(counting_date)
                    self._persist_day()
                    continue
                if self.current_state['counting_date'] != counting_date:
                    self.log('Daily cutoff reached - finalizing day totals')
                    self._persist_day()
                    self.current_state = None
                    self.save_state()
                    self._start_new_day(counting_date)
                    self._persist_day()

    def start_cutoff_watcher(self):
        t = threading.Thread(target=self.cutoff_watcher_loop, daemon=True)
        t.start()
        return t

    @property
    def current_count(self):
        if self.current_state is None:
            return 0
        return self.current_state['total_count']

    @property
    def current_counts(self):
        if self.current_state is None:
            return self._empty_counts()
        return dict(self.current_state.get('counts') or self._empty_counts())

    def _day_zone_counts(self, counting_date=None):
        if counting_date is None:
            counting_date = self.get_counting_date()
        cur = self.db.cursor()
        cur.execute(
            """
            SELECT zone_id, COALESCE(count, 0)
            FROM daily_zone_counts
            WHERE counting_date = ? AND camera_name = ? AND object_label = ?
            """,
            (counting_date, self.camera_name, self.object_label),
        )
        out = self._empty_counts()
        for zone_id, count in cur.fetchall():
            if zone_id in out:
                out[zone_id] = int(count)
        return out

    def display_total(self):
        cur = self.db.cursor()
        cur.execute(
            """
            SELECT COALESCE(total_count, 0)
            FROM daily_counters
            WHERE counting_date = ? AND camera_name = ? AND object_label = ?
            """,
            (self.get_counting_date(), self.camera_name, self.object_label),
        )
        row = cur.fetchone()
        return int(row[0]) if row else 0

    def display_counts(self):
        return self._day_zone_counts()

    def shutdown(self):
        self.shutdown_event.set()
        with self.state_lock:
            if self.current_state is not None:
                self._persist_day()
        self.db.close()