from __future__ import annotations

import json
import sqlite3
from datetime import datetime, timezone
from typing import Any


class JobsRepository:
    def __init__(self, conn: sqlite3.Connection) -> None:
        self._conn = conn

    @staticmethod
    def _now() -> str:
        return datetime.now(timezone.utc).isoformat(timespec="milliseconds").replace(
            "+00:00",
            "Z",
        )

    def create_job(
        self,
        job_id: str,
        method_id: str,
        method_version: int,
        status: str,
        hash_file_name: str,
        hash_file_sha256: str,
        hash_count: int,
        step_count: int,
        client_id: str | None = None,
    ) -> str:
        created_at = self._now()
        queued_at = created_at if status == "QUEUED" else None

        self._conn.execute(
            """
            INSERT INTO crack_jobs (
                id,
                method_id,
                method_version,
                status,
                created_at,
                queued_at,
                started_at,
                completed_at,
                cancelled_at,
                imported_at,
                client_id,
                hash_file_name,
                hash_file_sha256,
                hash_count,
                step_count,
                result_sha256,
                error
            )
            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            """,
            (
                job_id,
                method_id,
                method_version,
                status,
                created_at,
                queued_at,
                None,
                None,
                None,
                None,
                client_id,
                hash_file_name,
                hash_file_sha256,
                hash_count,
                step_count,
                None,
                None,
            ),
        )

        return created_at

    def get_job(self, job_id: str) -> sqlite3.Row | None:
        return self._conn.execute(
            """
            SELECT
                id,
                method_id,
                method_version,
                status,
                created_at,
                queued_at,
                started_at,
                completed_at,
                cancelled_at,
                imported_at,
                client_id,
                hash_file_name,
                hash_file_sha256,
                hash_count,
                step_count,
                result_sha256,
                error
            FROM crack_jobs
            WHERE id = ?
            """,
            (job_id,),
        ).fetchone()

    def list_jobs(
        self,
        status: str | None = None,
        method_id: str | None = None,
        client_id: str | None = None,
    ) -> list[sqlite3.Row]:
        query = """
            SELECT
                id,
                method_id,
                method_version,
                status,
                created_at,
                queued_at,
                started_at,
                completed_at,
                cancelled_at,
                imported_at,
                client_id,
                hash_file_name,
                hash_file_sha256,
                hash_count,
                step_count,
                result_sha256,
                error
            FROM crack_jobs
            WHERE 1 = 1
        """
        params: list[Any] = []

        if status is not None:
            query += "\n              AND status = ?"
            params.append(status)

        if method_id is not None:
            query += "\n              AND method_id = ?"
            params.append(method_id)

        if client_id is not None:
            query += "\n              AND client_id = ?"
            params.append(client_id)

        query += "\n            ORDER BY created_at DESC, id DESC"

        return self._conn.execute(query, params).fetchall()

    def update_job_status(
        self,
        job_id: str,
        status: str,
    ) -> None:
        now = self._now()

        timestamp_column = {
            "QUEUED": "queued_at",
            "RUNNING": "started_at",
            "COMPLETED": "completed_at",
            "CANCELLED": "cancelled_at",
        }.get(status)

        if timestamp_column is None:
            self._conn.execute(
                """
                UPDATE crack_jobs
                SET status = ?
                WHERE id = ?
                """,
                (
                    status,
                    job_id,
                ),
            )
            return

        self._conn.execute(
            f"""
            UPDATE crack_jobs
            SET status = ?,
                {timestamp_column} = ?
            WHERE id = ?
            """,
            (
                status,
                now,
                job_id,
            ),
        )

    def assign_client(
        self,
        job_id: str,
        client_id: str | None,
    ) -> None:
        self._conn.execute(
            """
            UPDATE crack_jobs
            SET client_id = ?
            WHERE id = ?
            """,
            (
                client_id,
                job_id,
            ),
        )

    def update_job_result(
        self,
        job_id: str,
        result_sha256: str | None = None,
        error: str | None = None,
        imported_at: str | None = None,
    ) -> None:
        self._conn.execute(
            """
            UPDATE crack_jobs
            SET result_sha256 = ?,
                error = ?,
                imported_at = ?
            WHERE id = ?
            """,
            (
                result_sha256,
                error,
                imported_at,
                job_id,
            ),
        )

    def add_job_handshake(
        self,
        job_id: str,
        handshake_id: int,
        hash22000: str,
        access_point_id: int | None = None,
    ) -> None:
        self._conn.execute(
            """
            INSERT INTO crack_job_handshakes (
                job_id,
                handshake_id,
                hash22000,
                access_point_id
            )
            VALUES (?, ?, ?, ?)
            """,
            (
                job_id,
                handshake_id,
                hash22000,
                access_point_id,
            ),
        )

    def get_job_handshake(
        self,
        job_id: str,
        handshake_id: int,
    ) -> sqlite3.Row | None:
        return self._conn.execute(
            """
            SELECT
                job_id,
                handshake_id,
                hash22000,
                access_point_id
            FROM crack_job_handshakes
            WHERE job_id = ?
              AND handshake_id = ?
            """,
            (
                job_id,
                handshake_id,
            ),
        ).fetchone()

    def list_job_handshakes(
        self,
        job_id: str,
    ) -> list[sqlite3.Row]:
        return self._conn.execute(
            """
            SELECT
                job_id,
                handshake_id,
                hash22000,
                access_point_id
            FROM crack_job_handshakes
            WHERE job_id = ?
            ORDER BY handshake_id
            """,
            (job_id,),
        ).fetchall()

    def create_job_step(
        self,
        job_id: str,
        step_no: int,
        step_id: str,
        status: str,
        session_name: str,
        definition: dict[str, Any],
    ) -> None:
        self._conn.execute(
            """
            INSERT INTO crack_job_steps (
                job_id,
                step_no,
                step_id,
                status,
                session_name,
                definition_json,
                started_at,
                completed_at,
                exit_code,
                restore_seen,
                error
            )
            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            """,
            (
                job_id,
                step_no,
                step_id,
                status,
                session_name,
                json.dumps(
                    definition,
                    ensure_ascii=False,
                    separators=(",", ":"),
                    sort_keys=True,
                ),
                None,
                None,
                None,
                0,
                None,
            ),
        )

    def get_job_step(
        self,
        job_id: str,
        step_no: int,
    ) -> sqlite3.Row | None:
        return self._conn.execute(
            """
            SELECT
                job_id,
                step_no,
                step_id,
                status,
                session_name,
                definition_json,
                started_at,
                completed_at,
                exit_code,
                restore_seen,
                error
            FROM crack_job_steps
            WHERE job_id = ?
              AND step_no = ?
            """,
            (
                job_id,
                step_no,
            ),
        ).fetchone()

    def list_job_steps(
        self,
        job_id: str,
    ) -> list[sqlite3.Row]:
        return self._conn.execute(
            """
            SELECT
                job_id,
                step_no,
                step_id,
                status,
                session_name,
                definition_json,
                started_at,
                completed_at,
                exit_code,
                restore_seen,
                error
            FROM crack_job_steps
            WHERE job_id = ?
            ORDER BY step_no
            """,
            (job_id,),
        ).fetchall()

    def update_job_step(
        self,
        job_id: str,
        step_no: int,
        status: str | None = None,
        started_at: str | None = None,
        completed_at: str | None = None,
        exit_code: int | None = None,
        restore_seen: bool | None = None,
        error: str | None = None,
    ) -> None:
        fields: list[str] = []
        params: list[Any] = []

        if status is not None:
            fields.append("status = ?")
            params.append(status)

        if started_at is not None:
            fields.append("started_at = ?")
            params.append(started_at)

        if completed_at is not None:
            fields.append("completed_at = ?")
            params.append(completed_at)

        if exit_code is not None:
            fields.append("exit_code = ?")
            params.append(exit_code)

        if restore_seen is not None:
            fields.append("restore_seen = ?")
            params.append(int(restore_seen))

        if error is not None:
            fields.append("error = ?")
            params.append(error)

        if not fields:
            return

        params.extend((job_id, step_no))

        self._conn.execute(
            f"""
            UPDATE crack_job_steps
            SET {", ".join(fields)}
            WHERE job_id = ?
              AND step_no = ?
            """,
            params,
        )