from __future__ import annotations

from dataclasses import dataclass
from typing import Any, Mapping, Optional

from google.cloud import bigquery

from .errors import CostExceeded, InvalidParameters
from .guardrails import enforce_guardrails
from .logging import get_logger
from .models import DatasetInfo, QueryConfig, QueryResult, TableInfo, TableSchema, FieldSchema

log = get_logger()


def _to_bq_params(params: Mapping[str, Any]) -> list[bigquery.ScalarQueryParameter]:
    out: list[bigquery.ScalarQueryParameter] = []
    for k, v in params.items():
        # Let BigQuery infer type for common python primitives by using STRING for str,
        # INT64 for int, FLOAT64 for float, BOOL for bool. Dates should be passed as ISO strings.
        if isinstance(v, bool):
            t = "BOOL"
        elif isinstance(v, int):
            t = "INT64"
        elif isinstance(v, float):
            t = "FLOAT64"
        else:
            t = "STRING"
        out.append(bigquery.ScalarQueryParameter(k, t, v))
    return out


@dataclass
class AgenticBQ:
    project: Optional[str] = None
    location: Optional[str] = None

    def __post_init__(self) -> None:
        self.client = bigquery.Client(project=self.project, location=self.location)

    # ---------- Querying ----------

    def dry_run_cost(
        self,
        sql: str,
        params: Optional[Mapping[str, Any]] = None,
        config: QueryConfig = QueryConfig(),
    ) -> int:
        guarded = enforce_guardrails(sql, config.deny_patterns, config.max_rows)

        bq_params = _to_bq_params(params or {})
        job_config = bigquery.QueryJobConfig(query_parameters=bq_params, dry_run=True, use_query_cache=False)

        job = self.client.query(guarded, job_config=job_config, job_id_prefix="agentic_dry_")
        # In dry-run, no rows are returned; metadata is populated.
        bytes_processed = int(job.total_bytes_processed or 0)
        log.info("Dry-run bytes_processed=%s", bytes_processed)
        return bytes_processed

    def safe_query(
        self,
        sql: str,
        params: Optional[Mapping[str, Any]] = None,
        config: QueryConfig = QueryConfig(),
    ) -> QueryResult:
        params = params or {}

        if config.require_params:
            # naive rule: if query uses @param but params dict is empty => reject
            if "@" in sql and not params:
                raise InvalidParameters("Query contains @params but params were not provided.")

        guarded = enforce_guardrails(sql, config.deny_patterns, config.max_rows)

        # Cost check
        bytes_processed = self.dry_run_cost(guarded, params=params, config=config)
        if bytes_processed > config.max_bytes_processed:
            raise CostExceeded(
                f"Dry-run cost exceeded: bytes_processed={bytes_processed} > max_bytes_processed={config.max_bytes_processed}"
            )

        bq_params = _to_bq_params(params)
        job_config = bigquery.QueryJobConfig(query_parameters=bq_params)

        job = self.client.query(guarded, job_config=job_config, job_id_prefix="agentic_")
        rows_iter = job.result()
        rows = [dict(r) for r in rows_iter]

        return QueryResult(
            sql=guarded,
            job_id=str(job.job_id),
            total_bytes_processed=int(job.total_bytes_processed or bytes_processed),
            rows=rows,
        )

    # ---------- Metadata ----------

    def list_datasets(self) -> list[DatasetInfo]:
        out: list[DatasetInfo] = []
        for ds in self.client.list_datasets():
            # ds is DatasetListItem
            out.append(DatasetInfo(project=ds.project, dataset=ds.dataset_id, location=getattr(ds, "location", None)))
        return out

    def list_tables(self, dataset: str, project: Optional[str] = None) -> list[TableInfo]:
        project = project or self.client.project
        ds_ref = bigquery.DatasetReference(project, dataset)
        out: list[TableInfo] = []
        for t in self.client.list_tables(ds_ref):
            out.append(TableInfo(project=t.project, dataset=t.dataset_id, table=t.table_id, table_type=t.table_type))
        return out

    def table_schema(self, dataset: str, table: str, project: Optional[str] = None) -> TableSchema:
        project = project or self.client.project
        table_ref = bigquery.TableReference(bigquery.DatasetReference(project, dataset), table)
        tbl = self.client.get_table(table_ref)
        fields: list[FieldSchema] = []
        for f in tbl.schema:
            fields.append(FieldSchema(name=f.name, field_type=f.field_type, mode=f.mode, description=f.description))
        return TableSchema(project=project, dataset=dataset, table=table, fields=fields)

