"""AI Atlas initial schema — entities, aliases, identifiers, sources, connectors, documents, snapshots, claims, relations,
change events, benchmark results, prices, jobs, LLM jobs, review queue, stats.

Revision ID: 0001
"""
from __future__ import annotations

from alembic import op

revision = "0001"
down_revision = None
branch_labels = None
depends_on = None

SQL = r"""
create extension if not exists pg_trgm;
create extension if not exists "uuid-ossp";
do $$ begin
  create extension if not exists vector;
exception when others then
  raise notice 'pgvector unavailable: %', sqlerrm;
end $$;

-- ---------------------------------------------------------------------------------------------- reference: sources & connectors
create table if not exists sources (
  id                text primary key,
  key               text not null unique,               -- e.g. 'openai.com', 'huggingface.co', 'arxiv.org'
  name              text not null,
  domain            text not null,
  organization_id   text,                               -- entity id of the owning organization (domain trust graph)
  tier              smallint not null default 2,        -- 1 official primary, 2 quality secondary, 3 community, 4 unverified
  kind              text not null default 'website',    -- website|docs|feed|repository|registry|dataset|leaderboard|regulator
  category          text not null default 'lab',
  base_url          text,
  robots_policy     text not null default 'respect',    -- respect|documented-exception
  rate_limit_per_min integer not null default 30,
  crawl_interval_s  integer not null default 86400,
  enabled           boolean not null default true,
  priority          smallint not null default 2,        -- 0 = P0
  notes             text,
  meta              jsonb not null default '{}'::jsonb,
  created_at        timestamptz not null default now(),
  updated_at        timestamptz not null default now()
);

create table if not exists connectors (
  name                 text primary key,
  source_id            text references sources(id),
  label                text not null,
  description          text,
  enabled              boolean not null default true,
  priority             smallint not null default 2,
  interval_seconds     integer not null default 3600,
  min_interval_seconds integer not null default 900,
  max_interval_seconds integer not null default 604800,
  parser_version       text not null default '1',
  rate_limit_per_min   integer not null default 30,
  owner                text not null default 'ai-atlas',
  expected_min_records integer not null default 0,      -- breakage detection: fewer records than this = suspected failure
  last_attempt_at      timestamptz,
  last_success_at      timestamptz,
  last_change_at       timestamptz,
  next_run_at          timestamptz,
  last_duration_ms     integer,
  consecutive_failures integer not null default 0,
  consecutive_unchanged integer not null default 0,
  circuit_open_until   timestamptz,
  health               text not null default 'unknown', -- ok|degraded|failing|disabled|unknown
  meta                 jsonb not null default '{}'::jsonb,
  created_at           timestamptz not null default now(),
  updated_at           timestamptz not null default now()
);

create table if not exists connector_runs (
  id               text primary key,
  connector_name   text not null references connectors(name),
  started_at       timestamptz not null,
  finished_at      timestamptz,
  status           text not null,                       -- running|success|unchanged|failed|skipped|suspect
  duration_ms      integer,
  docs_discovered  integer not null default 0,
  docs_fetched     integer not null default 0,
  docs_changed     integer not null default 0,
  docs_unchanged   integer not null default 0,
  docs_failed      integer not null default 0,
  entities_created integer not null default 0,
  entities_updated integer not null default 0,
  claims_written   integer not null default 0,
  relations_written integer not null default 0,
  events_emitted   integer not null default 0,
  error            text,
  meta             jsonb not null default '{}'::jsonb
);
create index if not exists connector_runs_name_idx on connector_runs (connector_name, started_at desc);

create table if not exists connector_errors (
  id             bigserial primary key,
  connector_name text not null,
  run_id         text,
  url            text,
  error_type     text not null,
  message        text not null,
  created_at     timestamptz not null default now()
);
create index if not exists connector_errors_name_idx on connector_errors (connector_name, created_at desc);

-- ---------------------------------------------------------------------------------------------- entities (unified graph)
create table if not exists entities (
  id              text primary key,
  entity_type     text not null,
  canonical_name  text not null,
  slug            text not null unique,
  description     text,
  status          text not null default 'active',     -- active|deprecated|retired|announced|unknown|merged
  organization_id text references entities(id),
  attributes      jsonb not null default '{}'::jsonb, -- current value per property (materialised from claims)
  provenance      jsonb not null default '{}'::jsonb, -- property -> {source, snapshot_id, url, observed_at, tier, confidence}
  quality         jsonb not null default '{}'::jsonb, -- source_count, primary_source_ratio, freshness, completeness, agreement, score
  counts          jsonb not null default '{}'::jsonb, -- cached relation / claim / event counts
  first_seen_at   timestamptz not null default now(),
  last_seen_at    timestamptz not null default now(),
  merged_into     text references entities(id),
  search          tsvector,
  created_at      timestamptz not null default now(),
  updated_at      timestamptz not null default now()
);
create index if not exists entities_type_idx on entities (entity_type, canonical_name);
create index if not exists entities_org_idx on entities (organization_id);
create index if not exists entities_search_idx on entities using gin (search);
create index if not exists entities_name_trgm_idx on entities using gin (canonical_name gin_trgm_ops);
create index if not exists entities_attrs_idx on entities using gin (attributes jsonb_path_ops);
create index if not exists entities_updated_idx on entities (updated_at desc);
create index if not exists entities_first_seen_idx on entities (first_seen_at desc);

create table if not exists entity_aliases (
  id          bigserial primary key,
  entity_id   text not null references entities(id) on delete cascade,
  alias       text not null,
  alias_norm  text not null,
  kind        text not null default 'alias',            -- alias|former_name|abbreviation|localized|typo
  snapshot_id text,
  created_at  timestamptz not null default now(),
  unique (entity_id, alias_norm)
);
create index if not exists entity_aliases_norm_idx on entity_aliases (alias_norm);

create table if not exists entity_identifiers (
  id          bigserial primary key,
  entity_id   text not null references entities(id) on delete cascade,
  scheme      text not null,                            -- hf_repo|github_repo|arxiv|doi|openrouter|domain|pypi|wikidata|url|provider_model_id
  value       text not null,
  snapshot_id text,
  created_at  timestamptz not null default now(),
  unique (scheme, value)
);
create index if not exists entity_identifiers_entity_idx on entity_identifiers (entity_id);

-- ---------------------------------------------------------------------------------------------- documents & snapshots (raw archive)
create table if not exists documents (
  id               text primary key,
  source_id        text references sources(id),
  connector_name   text references connectors(name),
  url              text not null unique,
  canonical_url    text,
  doc_type         text not null default 'page',       -- page|feed|feed_item|model_card|docs|pricing|release|paper|pdf|json|sitemap|repo|readme|leaderboard
  title            text,
  entity_id        text references entities(id),
  status           text not null default 'active',     -- active|gone|blocked|error
  first_seen_at    timestamptz not null default now(),
  last_fetched_at  timestamptz,
  last_changed_at  timestamptz,
  last_status      integer,
  etag             text,
  last_modified    text,
  content_hash     text,
  fetch_count      integer not null default 0,
  change_count     integer not null default 0,
  fail_count       integer not null default 0,
  fetch_interval_s integer,
  next_fetch_at    timestamptz,
  priority         smallint not null default 2,
  needs_llm        boolean not null default false,
  meta             jsonb not null default '{}'::jsonb
);
create index if not exists documents_connector_idx on documents (connector_name, last_fetched_at desc);
create index if not exists documents_entity_idx on documents (entity_id);
create index if not exists documents_next_idx on documents (next_fetch_at) where status = 'active';

create table if not exists snapshots (
  id                text primary key,
  document_id       text not null references documents(id) on delete cascade,
  run_id            text,
  url               text not null,
  final_url         text,
  observed_at       timestamptz not null default now(),
  http_status       integer,
  headers           jsonb not null default '{}'::jsonb,
  content_type      text,
  content_hash      text not null,
  byte_size         integer not null default 0,
  raw_path          text,                              -- content-addressed gzip under AIA_DATA_DIR/raw
  text_path         text,                              -- cleaned text (gzip) under AIA_DATA_DIR/text
  text_hash         text,
  structured        jsonb,                             -- deterministic extraction (json-ld, meta, tables, embedded json)
  parser_version    text not null default '1',
  connector_version text not null default '1',
  transport         text not null default 'direct',    -- direct|browser|scrapfly|firecrawl|file|git
  changed           boolean not null default true,
  diff              jsonb,                             -- structural diff vs previous changed snapshot
  processing_status text not null default 'stored',    -- stored|extracted|llm_pending|llm_done|failed
  created_at        timestamptz not null default now()
);
create index if not exists snapshots_doc_idx on snapshots (document_id, observed_at desc);
create index if not exists snapshots_hash_idx on snapshots (content_hash);
create index if not exists snapshots_status_idx on snapshots (processing_status) where processing_status in ('stored','llm_pending');

-- ---------------------------------------------------------------------------------------------- temporal facts
create table if not exists claims (
  id                text primary key,
  entity_id         text not null references entities(id) on delete cascade,
  property          text not null,
  value             jsonb not null,
  value_text        text,
  value_num         double precision,
  unit              text,
  source_id         text references sources(id),
  snapshot_id       text references snapshots(id),
  source_url        text,
  tier              smallint not null default 2,
  confidence        text not null default 'medium',     -- verified|high|medium|low|conflicted
  status            text not null default 'current',    -- current|superseded|conflicting|retracted
  extractor         text not null default 'deterministic',
  extractor_version text not null default '1',
  observed_at       timestamptz not null default now(),
  effective_at      timestamptz,
  valid_from        timestamptz not null default now(),
  valid_to          timestamptz,
  created_at        timestamptz not null default now()
);
create index if not exists claims_entity_prop_idx on claims (entity_id, property, valid_from desc);
create index if not exists claims_current_idx on claims (entity_id, property) where status = 'current';
create index if not exists claims_snapshot_idx on claims (snapshot_id);
create index if not exists claims_observed_idx on claims (observed_at desc);

create table if not exists relations (
  id          text primary key,
  subject_id  text not null references entities(id) on delete cascade,
  predicate   text not null,
  object_id   text not null references entities(id) on delete cascade,
  attributes  jsonb not null default '{}'::jsonb,
  source_id   text references sources(id),
  snapshot_id text references snapshots(id),
  source_url  text,
  tier        smallint not null default 2,
  confidence  text not null default 'medium',
  observed_at timestamptz not null default now(),
  valid_from  timestamptz not null default now(),
  valid_to    timestamptz,
  created_at  timestamptz not null default now()
);
create unique index if not exists relations_live_uniq on relations (subject_id, predicate, object_id) where valid_to is null;
create index if not exists relations_subject_idx on relations (subject_id, predicate);
create index if not exists relations_object_idx on relations (object_id, predicate);

create table if not exists change_events (
  id             text primary key,
  entity_id      text references entities(id) on delete cascade,
  event_type     text not null,                          -- NEW_MODEL|MODEL_UPDATED|PRICE_CHANGED|CONTEXT_CHANGED|NEW_PAPER|RELEASE|BENCHMARK_RESULT|NEW_PROVIDER|DOCUMENT_CHANGED|…
  category       text not null default 'update',         -- model|price|benchmark|paper|release|company|hardware|framework|provider|dataset|regulation|incident|repository
  property       text,
  old_value      jsonb,
  new_value      jsonb,
  summary        text,
  importance     smallint not null default 2,            -- 0 minor … 3 major
  observed_at    timestamptz not null default now(),
  effective_at   timestamptz,
  source_id      text references sources(id),
  snapshot_id    text references snapshots(id),
  source_url     text,
  connector_name text,
  dedupe_key     text unique,
  meta           jsonb not null default '{}'::jsonb
);
create index if not exists change_events_observed_idx on change_events (observed_at desc);
create index if not exists change_events_entity_idx on change_events (entity_id, observed_at desc);
create index if not exists change_events_type_idx on change_events (event_type, observed_at desc);
create index if not exists change_events_category_idx on change_events (category, observed_at desc);

-- ---------------------------------------------------------------------------------------------- domain tables
create table if not exists benchmark_results (
  id            text primary key,
  model_id      text not null references entities(id) on delete cascade,
  benchmark_id  text not null references entities(id) on delete cascade,
  score         double precision not null,
  metric        text,
  unit          text,
  higher_is_better boolean not null default true,
  config        jsonb not null default '{}'::jsonb,     -- prompting, judge, tool use, sampling, harness, model variant string
  evaluated_at  timestamptz,
  observed_at   timestamptz not null default now(),
  source_id     text references sources(id),
  snapshot_id   text references snapshots(id),
  source_url    text,
  tier          smallint not null default 2,
  confidence    text not null default 'medium',
  dedupe_key    text unique,
  valid_to      timestamptz
);
create index if not exists benchmark_results_model_idx on benchmark_results (model_id, benchmark_id, observed_at desc);
create index if not exists benchmark_results_bench_idx on benchmark_results (benchmark_id, score desc);

create table if not exists prices (
  id                     text primary key,
  model_id               text not null references entities(id) on delete cascade,
  provider_id            text not null references entities(id) on delete cascade,
  provider_model_id      text,
  input_per_mtok         double precision,
  output_per_mtok        double precision,
  cached_input_per_mtok  double precision,
  cache_write_per_mtok   double precision,
  batch_input_per_mtok   double precision,
  batch_output_per_mtok  double precision,
  per_image              double precision,
  per_request            double precision,
  currency               text not null default 'USD',
  context_length         integer,
  max_output_tokens      integer,
  features               jsonb not null default '{}'::jsonb,
  observed_at            timestamptz not null default now(),
  valid_from             timestamptz not null default now(),
  valid_to               timestamptz,
  source_id              text references sources(id),
  snapshot_id            text references snapshots(id),
  source_url             text,
  tier                   smallint not null default 2,
  meta                   jsonb not null default '{}'::jsonb
);
create unique index if not exists prices_live_uniq on prices (model_id, provider_id, coalesce(provider_model_id, '')) where valid_to is null;
create index if not exists prices_provider_idx on prices (provider_id, valid_from desc);
create index if not exists prices_model_idx on prices (model_id, valid_from desc);

create table if not exists domains (
  domain          text primary key,
  organization_id text references entities(id),
  trust_tier      smallint not null default 3,
  category        text,
  notes           text,
  created_at      timestamptz not null default now()
);

-- ---------------------------------------------------------------------------------------------- work queues
create table if not exists jobs (
  id           text primary key,
  kind         text not null,                            -- fetch_document|llm_extract|embed_entity|reprocess_snapshot|recompute_quality|resolve_entity
  payload      jsonb not null default '{}'::jsonb,
  priority     smallint not null default 5,              -- 0 highest
  status       text not null default 'queued',           -- queued|running|done|failed|dead
  attempts     integer not null default 0,
  max_attempts integer not null default 3,
  run_after    timestamptz not null default now(),
  locked_by    text,
  locked_at    timestamptz,
  started_at   timestamptz,
  finished_at  timestamptz,
  error        text,
  batch_id     text,
  dedupe_key   text,
  created_at   timestamptz not null default now()
);
create index if not exists jobs_pick_idx on jobs (status, priority, run_after) where status = 'queued';
create unique index if not exists jobs_dedupe_idx on jobs (dedupe_key) where status in ('queued','running');
create index if not exists jobs_batch_idx on jobs (batch_id);

create table if not exists llm_jobs (
  id             text primary key,
  job_id         text,
  task_type      text not null,
  stage          text not null,                          -- small|medium|large
  engine         text not null,
  model          text not null,
  node           text,
  schema_name    text,
  snapshot_id    text references snapshots(id),
  entity_id      text references entities(id),
  input_tokens   integer,
  output_tokens  integer,
  duration_ms    integer,
  status         text not null,                          -- ok|failed|invalid_json|schema_error
  output         jsonb,
  error          text,
  created_at     timestamptz not null default now()
);
create index if not exists llm_jobs_created_idx on llm_jobs (created_at desc);

create table if not exists review_queue (
  id          text primary key,
  kind        text not null,                             -- merge_candidate|conflict|blocked_source|parser_breakage|unusual_change|new_source
  entity_ids  text[] not null default '{}',
  payload     jsonb not null default '{}'::jsonb,
  reason      text not null,
  status      text not null default 'pending',           -- pending|approved|rejected|edited
  resolution  jsonb,
  created_at  timestamptz not null default now(),
  resolved_at timestamptz,
  dedupe_key  text unique
);
create index if not exists review_queue_status_idx on review_queue (status, created_at desc);

-- ---------------------------------------------------------------------------------------------- embeddings (optional; local models)
do $$ begin
  if exists (select 1 from pg_extension where extname = 'vector') then
    execute 'create table if not exists entity_embeddings (
      entity_id text primary key references entities(id) on delete cascade,
      model text not null,
      embedding vector(1024) not null,
      text_hash text not null,
      created_at timestamptz not null default now())';
  end if;
end $$;

-- ---------------------------------------------------------------------------------------------- stats, views, api keys
create table if not exists stats_snapshots (
  id          bigserial primary key,
  computed_at timestamptz not null default now(),
  counts      jsonb not null
);

create table if not exists page_views (
  path  text not null,
  day   date not null,
  views integer not null default 0,
  primary key (path, day)
);

create table if not exists api_keys (
  id          text primary key,
  key_hash    text not null unique,
  label       text not null,
  owner_email text,
  plan        text not null default 'developer',
  rate_per_min integer not null default 60,
  enabled     boolean not null default true,
  created_at  timestamptz not null default now(),
  last_used_at timestamptz,
  usage_count bigint not null default 0
);

create table if not exists metric_definitions (
  key         text primary key,
  label       text not null,
  version     text not null,
  description text not null,
  formula     text,
  created_at  timestamptz not null default now()
);

-- ---------------------------------------------------------------------------------------------- search vector maintenance
create or replace function entities_search_update() returns trigger language plpgsql as $$
begin
  new.search :=
    setweight(to_tsvector('simple', coalesce(new.canonical_name, '')), 'A') ||
    setweight(to_tsvector('simple', coalesce(new.attributes->>'family', '')), 'B') ||
    setweight(to_tsvector('simple', coalesce(new.entity_type, '')), 'C') ||
    setweight(to_tsvector('english', left(coalesce(new.description, ''), 4000)), 'C');
  new.updated_at := now();
  return new;
end $$;
drop trigger if exists entities_search_trg on entities;
create trigger entities_search_trg before insert or update of canonical_name, description, attributes, entity_type on entities
  for each row execute function entities_search_update();
"""


def upgrade() -> None:
    for statement in _split(SQL):
        op.execute(statement)


def downgrade() -> None:
    raise RuntimeError("AI Atlas migrations are forward-only: historical data is never disposable")


def _split(sql: str) -> list[str]:
    """Split on semicolons outside `$$ … $$` blocks."""
    out: list[str] = []
    buf: list[str] = []
    in_dollar = False
    for line in sql.splitlines():
        stripped = line.strip()
        if stripped.count("$$") % 2 == 1:
            in_dollar = not in_dollar
        buf.append(line)
        if not in_dollar and stripped.endswith(";"):
            body = [ln for ln in buf if not ln.strip().startswith("--") or in_dollar]
            stmt = "\n".join(body).strip()
            if stmt:
                out.append(stmt)
            buf = []
    tail = "\n".join(ln for ln in buf if not ln.strip().startswith("--")).strip()
    if tail:
        out.append(tail)
    return out