from __future__ import annotations

import json
from datetime import UTC, datetime
from typing import Any

from google.cloud import bigquery

from .naming import bq_path


class LeaseUnavailableError(RuntimeError):
    pass


class BigQueryOutboundState:
    def __init__(self, client: bigquery.Client, *, project: str, dataset: str, location: str) -> None:
        self.client = client
        self.project = project
        self.dataset = dataset
        self.location = location

    @property
    def leases_table(self) -> str:
        return f"{self.project}.{self.dataset}.outbound_leases"

    @property
    def runs_table(self) -> str:
        return f"{self.project}.{self.dataset}.outbound_runs"

    @property
    def attempts_table(self) -> str:
        return f"{self.project}.{self.dataset}.outbound_record_attempts"

    def ensure_tables(self) -> None:
        self._query(f"""
CREATE TABLE IF NOT EXISTS {bq_path(self.project, self.dataset, "outbound_leases")} (
  job_name STRING NOT NULL,
  holder_run_id STRING,
  acquired_at TIMESTAMP,
  expires_at TIMESTAMP,
  updated_at TIMESTAMP
)
""")
        self._query(f"""
CREATE TABLE IF NOT EXISTS {bq_path(self.project, self.dataset, "outbound_runs")} (
  job_name STRING NOT NULL,
  run_id STRING NOT NULL,
  status STRING NOT NULL,
  dry_run BOOL,
  started_at TIMESTAMP,
  ended_at TIMESTAMP,
  metrics_json STRING,
  error_message STRING
)
PARTITION BY DATE(started_at)
CLUSTER BY job_name, status
""")
        self._query(f"""
CREATE TABLE IF NOT EXISTS {bq_path(self.project, self.dataset, "outbound_record_attempts")} (
  job_name STRING NOT NULL,
  run_id STRING NOT NULL,
  external_key STRING NOT NULL,
  source_hash STRING NOT NULL,
  target_id STRING,
  status STRING NOT NULL,
  created BOOL,
  error_code STRING,
  error_message STRING,
  attempted_at TIMESTAMP NOT NULL
)
PARTITION BY DATE(attempted_at)
CLUSTER BY job_name, external_key, status
""")

    def acquire_lease(self, job_name: str, run_id: str, lease_seconds: int) -> None:
        params = [
            bigquery.ScalarQueryParameter("job_name", "STRING", job_name),
            bigquery.ScalarQueryParameter("run_id", "STRING", run_id),
        ]
        self._query(
            f"""
MERGE {bq_path(self.project, self.dataset, "outbound_leases")} T
USING (
  SELECT @job_name AS job_name, @run_id AS run_id,
         TIMESTAMP_ADD(CURRENT_TIMESTAMP(), INTERVAL {int(lease_seconds)} SECOND) AS expires_at
) S
ON T.job_name = S.job_name
WHEN MATCHED AND (T.expires_at IS NULL OR T.expires_at <= CURRENT_TIMESTAMP() OR T.holder_run_id = S.run_id)
  THEN UPDATE SET holder_run_id = S.run_id, acquired_at = CURRENT_TIMESTAMP(),
                  expires_at = S.expires_at, updated_at = CURRENT_TIMESTAMP()
WHEN NOT MATCHED
  THEN INSERT (job_name, holder_run_id, acquired_at, expires_at, updated_at)
       VALUES (S.job_name, S.run_id, CURRENT_TIMESTAMP(), S.expires_at, CURRENT_TIMESTAMP())
""",
            params,
        )
        rows = list(
            self._query(
                f"""
SELECT holder_run_id, expires_at
FROM {bq_path(self.project, self.dataset, "outbound_leases")}
WHERE job_name = @job_name
""",
                [bigquery.ScalarQueryParameter("job_name", "STRING", job_name)],
            )
        )
        if len(rows) != 1 or str(rows[0]["holder_run_id"]) != run_id:
            holder = str(rows[0]["holder_run_id"]) if rows else "unknown"
            raise LeaseUnavailableError(f"outbound job {job_name} is already held by run {holder}")

    def release_lease(self, job_name: str, run_id: str) -> None:
        self._query(
            f"""
UPDATE {bq_path(self.project, self.dataset, "outbound_leases")}
SET holder_run_id = NULL, expires_at = CURRENT_TIMESTAMP(), updated_at = CURRENT_TIMESTAMP()
WHERE job_name = @job_name AND holder_run_id = @run_id
""",
            [
                bigquery.ScalarQueryParameter("job_name", "STRING", job_name),
                bigquery.ScalarQueryParameter("run_id", "STRING", run_id),
            ],
        )

    def start_run(self, job_name: str, run_id: str) -> None:
        self._query(
            f"""
INSERT INTO {bq_path(self.project, self.dataset, "outbound_runs")}
  (job_name, run_id, status, dry_run, started_at)
VALUES (@job_name, @run_id, 'running', FALSE, CURRENT_TIMESTAMP())
""",
            [
                bigquery.ScalarQueryParameter("job_name", "STRING", job_name),
                bigquery.ScalarQueryParameter("run_id", "STRING", run_id),
            ],
        )

    def finish_run(
        self,
        job_name: str,
        run_id: str,
        *,
        status: str,
        metrics: dict[str, Any],
        error_message: str | None = None,
    ) -> None:
        self._query(
            f"""
UPDATE {bq_path(self.project, self.dataset, "outbound_runs")}
SET status = @status, ended_at = CURRENT_TIMESTAMP(), metrics_json = @metrics_json,
    error_message = @error_message
WHERE job_name = @job_name AND run_id = @run_id
""",
            [
                bigquery.ScalarQueryParameter("status", "STRING", status),
                bigquery.ScalarQueryParameter("metrics_json", "STRING", json.dumps(metrics, sort_keys=True)),
                bigquery.ScalarQueryParameter("error_message", "STRING", error_message),
                bigquery.ScalarQueryParameter("job_name", "STRING", job_name),
                bigquery.ScalarQueryParameter("run_id", "STRING", run_id),
            ],
        )

    def latest_record_states(self, job_name: str) -> dict[str, dict[str, Any]]:
        rows = self._query(
            f"""
SELECT external_key, source_hash, target_id, status
FROM {bq_path(self.project, self.dataset, "outbound_record_attempts")}
WHERE job_name = @job_name
QUALIFY ROW_NUMBER() OVER (PARTITION BY external_key ORDER BY attempted_at DESC, run_id DESC) = 1
""",
            [bigquery.ScalarQueryParameter("job_name", "STRING", job_name)],
        )
        return {
            str(row["external_key"]): {
                "source_hash": str(row["source_hash"]),
                "target_id": str(row["target_id"] or ""),
                "status": str(row["status"]),
            }
            for row in rows
        }

    def append_attempts(self, rows: list[dict[str, Any]]) -> None:
        if not rows:
            return
        now = datetime.now(UTC).isoformat()
        payload = [{**row, "attempted_at": now} for row in rows]
        errors = self.client.insert_rows_json(self.attempts_table, payload)
        if errors:
            raise RuntimeError(f"failed to persist outbound attempt ledger: {errors}")

    def _query(
        self, sql: str, params: list[bigquery.ScalarQueryParameter] | None = None
    ) -> bigquery.table.RowIterator:
        job_config = bigquery.QueryJobConfig(query_parameters=params or [])
        return self.client.query(sql, job_config=job_config, location=self.location).result()