import sqlite3
import json
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional, Dict, Any, List
from agent.config import AGENT_HOME

STATE_DB_PATH = AGENT_HOME / "agent_state.db"

class StateDatabase:
    """Local SQLite database for agent offline resiliency, session resume tracking and file caching."""

    def __init__(self, db_path: Path = STATE_DB_PATH):
        self.db_path = db_path
        self._init_db()

    def _get_conn(self) -> sqlite3.Connection:
        return sqlite3.connect(str(self.db_path))

    def _init_db(self):
        with self._get_conn() as conn:
            cursor = conn.cursor()
            cursor.execute("""
                CREATE TABLE IF NOT EXISTS upload_sessions (
                    filepath TEXT PRIMARY KEY,
                    sha256 TEXT NOT NULL,
                    session_code TEXT NOT NULL,
                    total_chunks INTEGER NOT NULL,
                    chunk_size INTEGER NOT NULL,
                    completed_chunks INTEGER DEFAULT 0,
                    status TEXT NOT NULL,
                    last_updated TIMESTAMP NOT NULL
                )
            """)
            cursor.execute("""
                CREATE TABLE IF NOT EXISTS completed_files (
                    filepath TEXT PRIMARY KEY,
                    sha256 TEXT NOT NULL,
                    file_size INTEGER NOT NULL,
                    job_id INTEGER,
                    last_backup_time TIMESTAMP NOT NULL
                )
            """)
            conn.commit()

    def save_session(
        self,
        filepath: str,
        sha256: str,
        session_code: str,
        total_chunks: int,
        chunk_size: int,
        status: str = "UPLOADING"
    ):
        with self._get_conn() as conn:
            cursor = conn.cursor()
            now = datetime.now(timezone.utc).isoformat()
            cursor.execute("""
                INSERT OR REPLACE INTO upload_sessions 
                (filepath, sha256, session_code, total_chunks, chunk_size, status, last_updated)
                VALUES (?, ?, ?, ?, ?, ?, ?)
            """, (filepath, sha256, session_code, total_chunks, chunk_size, status, now))
            conn.commit()

    def get_session(self, filepath: str) -> Optional[Dict[str, Any]]:
        with self._get_conn() as conn:
            cursor = conn.cursor()
            cursor.execute("SELECT filepath, sha256, session_code, total_chunks, chunk_size, completed_chunks, status FROM upload_sessions WHERE filepath = ?", (filepath,))
            row = cursor.fetchone()
            if row:
                return {
                    "filepath": row[0],
                    "sha256": row[1],
                    "session_code": row[2],
                    "total_chunks": row[3],
                    "chunk_size": row[4],
                    "completed_chunks": row[5],
                    "status": row[6]
                }
            return None

    def mark_session_completed(self, filepath: str, sha256: str, file_size: int, job_id: Optional[int] = None):
        with self._get_conn() as conn:
            cursor = conn.cursor()
            now = datetime.now(timezone.utc).isoformat()
            cursor.execute("DELETE FROM upload_sessions WHERE filepath = ?", (filepath,))
            cursor.execute("""
                INSERT OR REPLACE INTO completed_files (filepath, sha256, file_size, job_id, last_backup_time)
                VALUES (?, ?, ?, ?, ?)
            """, (filepath, sha256, file_size, job_id, now))
            conn.commit()

    def is_file_already_backed_up(self, filepath: str, current_sha256: str) -> bool:
        with self._get_conn() as conn:
            cursor = conn.cursor()
            cursor.execute("SELECT sha256 FROM completed_files WHERE filepath = ?", (filepath,))
            row = cursor.fetchone()
            if row and row[0] == current_sha256:
                return True
            return False

state_db = StateDatabase()