"""
SQLite layer. Idempotente: cada corrida hace upsert por arxiv_id sin duplicar.

Decisiones de diseño:
- Listas (autores, categorias, topics, code_urls) se serializan a JSON.
  Para volumen <1M de papers, esto es mas simple que tablas hijas y
  mantiene la query principal en una sola tabla.
- Indices solo en columnas usadas para filtrar/ordenar: published, primary_cat, score.
- updated_at se setea automaticamente en cada upsert para detectar staleness.
"""

from __future__ import annotations

import json
import sqlite3
from contextlib import contextmanager
from datetime import datetime, date
from pathlib import Path
from typing import Any, Iterable, Iterator


SCHEMA = """
CREATE TABLE IF NOT EXISTS papers (
    arxiv_id        TEXT PRIMARY KEY,
    title           TEXT NOT NULL,
    abstract        TEXT,
    authors_json    TEXT,           -- JSON array of strings
    categories_json TEXT,           -- JSON array of strings
    primary_cat     TEXT,
    published       TEXT,           -- ISO date YYYY-MM-DD
    updated         TEXT,           -- ISO date YYYY-MM-DD
    doi             TEXT,
    pdf_url         TEXT,
    comment         TEXT,           -- arXiv "comments" field (puede traer github)
    journal_ref     TEXT,

    -- OpenAlex enrichment
    openalex_id        TEXT,
    cited_by_count     INTEGER,
    fwci               REAL,
    topics_json        TEXT,           -- JSON array
    publication_year   INTEGER,        -- ano de publicacion oficial (journal o arXiv)
    openalex_abstract  TEXT,           -- abstract reconstruido del inverted index

    -- Semantic Scholar enrichment (opcional)
    ss_id           TEXT,
    influential_citation_count INTEGER,

    -- Detected artifacts
    has_code        INTEGER,        -- 0/1
    has_data        INTEGER,        -- 0/1
    code_urls_json  TEXT,           -- JSON array
    data_urls_json  TEXT,           -- JSON array

    -- Ranking
    score           REAL,

    -- Metadata
    fetched_at      TEXT NOT NULL,  -- ISO datetime
    updated_at      TEXT NOT NULL   -- ISO datetime
);

CREATE INDEX IF NOT EXISTS idx_published    ON papers(published);
CREATE INDEX IF NOT EXISTS idx_primary_cat  ON papers(primary_cat);
CREATE INDEX IF NOT EXISTS idx_score        ON papers(score DESC);
CREATE INDEX IF NOT EXISTS idx_cited_by     ON papers(cited_by_count DESC);

-- Embeddings vectoriales para busqueda semantica.
-- Tabla separada (no en papers) para mantener exports CSV/Parquet lean.
CREATE TABLE IF NOT EXISTS embeddings (
    arxiv_id     TEXT PRIMARY KEY,
    model        TEXT NOT NULL,         -- nombre del modelo (ej. allenai-specter)
    dim          INTEGER NOT NULL,
    vector       BLOB NOT NULL,         -- np.float32 bytes, ya normalizado L2
    computed_at  TEXT NOT NULL,
    FOREIGN KEY (arxiv_id) REFERENCES papers(arxiv_id)
);
CREATE INDEX IF NOT EXISTS idx_emb_model ON embeddings(model);

-- Tabla auxiliar: log de corridas para auditoria y debugging
CREATE TABLE IF NOT EXISTS fetch_runs (
    run_id          INTEGER PRIMARY KEY AUTOINCREMENT,
    started_at      TEXT NOT NULL,
    finished_at     TEXT,
    category        TEXT,
    n_fetched       INTEGER,
    n_new           INTEGER,
    n_updated       INTEGER,
    status          TEXT,           -- 'ok' | 'partial' | 'error'
    error_msg       TEXT
);
"""


# ------------------------------------------------------------
# Connection helpers
# ------------------------------------------------------------

def connect(db_path: str | Path) -> sqlite3.Connection:
    """Abre conexion con pragmas razonables para uso single-writer."""
    db_path = Path(db_path)
    db_path.parent.mkdir(parents=True, exist_ok=True)
    conn = sqlite3.connect(str(db_path))
    conn.row_factory = sqlite3.Row
    # WAL = mejor concurrencia lectura/escritura, menos lock contention
    conn.execute("PRAGMA journal_mode = WAL")
    conn.execute("PRAGMA synchronous = NORMAL")
    conn.execute("PRAGMA foreign_keys = ON")
    return conn


def init_schema(conn: sqlite3.Connection) -> None:
    conn.executescript(SCHEMA)
    migrate_schema(conn)
    conn.commit()


def migrate_schema(conn: sqlite3.Connection) -> None:
    """
    Migraciones idempotentes para DBs creadas con versiones viejas del schema.
    SQLite no soporta `ALTER TABLE ADD COLUMN IF NOT EXISTS`, asi que
    consultamos `PRAGMA table_info` y agregamos solo lo que falta.
    """
    existing = {row[1] for row in conn.execute("PRAGMA table_info(papers)")}
    additions = {
        "publication_year": "INTEGER",
        "openalex_abstract": "TEXT",
        "user_rating": "INTEGER",           # -1 dislike, 0/NULL neutral, +1 like
        "user_note":   "TEXT",              # nota libre del usuario
        "rated_at":    "TEXT",              # ISO datetime del ultimo rating
        "methodology_json": "TEXT",         # JSON: [{"family","sem","lex"}] metodologias detectadas
    }
    for col, coltype in additions.items():
        if col not in existing:
            conn.execute(f"ALTER TABLE papers ADD COLUMN {col} {coltype}")
    # Indice para queries de "mi biblioteca" rapidas
    conn.execute("CREATE INDEX IF NOT EXISTS idx_user_rating ON papers(user_rating)")


@contextmanager
def session(db_path: str | Path) -> Iterator[sqlite3.Connection]:
    conn = connect(db_path)
    try:
        init_schema(conn)
        yield conn
        conn.commit()
    except Exception:
        conn.rollback()
        raise
    finally:
        conn.close()


# ------------------------------------------------------------
# Serialization helpers
# ------------------------------------------------------------

def _to_iso(d: Any) -> str | None:
    if d is None:
        return None
    if isinstance(d, (datetime, date)):
        return d.isoformat() if isinstance(d, datetime) else d.isoformat()
    return str(d)


def _jdump(x: Any) -> str | None:
    if x is None:
        return None
    return json.dumps(x, ensure_ascii=False)


def _jload(s: str | None) -> Any:
    if s is None or s == "":
        return None
    return json.loads(s)


# ------------------------------------------------------------
# Upserts
# ------------------------------------------------------------

# Campos provistos por el fetch inicial de arXiv (sin enrichment ni score).
ARXIV_FIELDS = [
    "arxiv_id", "title", "abstract", "authors_json", "categories_json",
    "primary_cat", "published", "updated", "doi", "pdf_url", "comment",
    "journal_ref",
]


def upsert_arxiv_paper(conn: sqlite3.Connection, paper: dict) -> str:
    """
    Inserta o actualiza un paper con datos crudos de arXiv.
    No toca campos de enrichment ni score.

    Returns: 'inserted' | 'updated'
    """
    now = datetime.utcnow().isoformat(timespec="seconds")
    row = {
        "arxiv_id":        paper["arxiv_id"],
        "title":           paper["title"],
        "abstract":        paper.get("abstract"),
        "authors_json":    _jdump(paper.get("authors") or []),
        "categories_json": _jdump(paper.get("categories") or []),
        "primary_cat":     paper.get("primary_cat"),
        "published":       _to_iso(paper.get("published")),
        "updated":         _to_iso(paper.get("updated")),
        "doi":             paper.get("doi"),
        "pdf_url":         paper.get("pdf_url"),
        "comment":         paper.get("comment"),
        "journal_ref":     paper.get("journal_ref"),
        "fetched_at":      now,
        "updated_at":      now,
    }

    cur = conn.execute(
        "SELECT 1 FROM papers WHERE arxiv_id = ?", (row["arxiv_id"],)
    )
    exists = cur.fetchone() is not None

    if exists:
        # UPDATE solo de campos de arXiv; preserva enrichment y score.
        conn.execute(
            """
            UPDATE papers SET
                title=:title, abstract=:abstract, authors_json=:authors_json,
                categories_json=:categories_json, primary_cat=:primary_cat,
                published=:published, updated=:updated, doi=:doi,
                pdf_url=:pdf_url, comment=:comment, journal_ref=:journal_ref,
                updated_at=:updated_at
            WHERE arxiv_id=:arxiv_id
            """,
            row,
        )
        return "updated"
    else:
        conn.execute(
            """
            INSERT INTO papers (
                arxiv_id, title, abstract, authors_json, categories_json,
                primary_cat, published, updated, doi, pdf_url, comment,
                journal_ref, fetched_at, updated_at
            ) VALUES (
                :arxiv_id, :title, :abstract, :authors_json, :categories_json,
                :primary_cat, :published, :updated, :doi, :pdf_url, :comment,
                :journal_ref, :fetched_at, :updated_at
            )
            """,
            row,
        )
        return "inserted"


def upsert_many(conn: sqlite3.Connection, papers: Iterable[dict]) -> dict:
    """Bulk upsert. Devuelve contadores."""
    counts = {"inserted": 0, "updated": 0}
    for p in papers:
        result = upsert_arxiv_paper(conn, p)
        counts[result] += 1
    conn.commit()
    return counts


# ------------------------------------------------------------
# Run log
# ------------------------------------------------------------

def start_run(conn: sqlite3.Connection, category: str) -> int:
    cur = conn.execute(
        "INSERT INTO fetch_runs (started_at, category, status) VALUES (?, ?, 'running')",
        (datetime.utcnow().isoformat(timespec="seconds"), category),
    )
    conn.commit()
    return cur.lastrowid


def finish_run(
    conn: sqlite3.Connection,
    run_id: int,
    n_fetched: int,
    n_new: int,
    n_updated: int,
    status: str = "ok",
    error_msg: str | None = None,
) -> None:
    conn.execute(
        """
        UPDATE fetch_runs SET
            finished_at = ?, n_fetched = ?, n_new = ?, n_updated = ?,
            status = ?, error_msg = ?
        WHERE run_id = ?
        """,
        (
            datetime.utcnow().isoformat(timespec="seconds"),
            n_fetched, n_new, n_updated, status, error_msg, run_id,
        ),
    )
    conn.commit()


# ------------------------------------------------------------
# Read helpers (utiles para los siguientes modulos)
# ------------------------------------------------------------

def count_papers(conn: sqlite3.Connection) -> int:
    return conn.execute("SELECT COUNT(*) FROM papers").fetchone()[0]


def papers_missing_openalex(conn: sqlite3.Connection, limit: int | None = None) -> list[sqlite3.Row]:
    sql = (
        "SELECT arxiv_id, doi, title FROM papers "
        "WHERE openalex_id IS NULL ORDER BY published DESC"
    )
    if limit:
        sql += f" LIMIT {int(limit)}"
    return conn.execute(sql).fetchall()


# ------------------------------------------------------------
# Enrichment upserts
# ------------------------------------------------------------

def update_openalex(
    conn: sqlite3.Connection,
    arxiv_id: str,
    openalex_id: str,
    cited_by_count: int | None,
    fwci: float | None,
    topics: list[dict] | None,
    publication_year: int | None = None,
    openalex_abstract: str | None = None,
) -> None:
    """Rellena los campos OpenAlex de un paper ya existente."""
    conn.execute(
        """
        UPDATE papers SET
            openalex_id = ?,
            cited_by_count = ?,
            fwci = ?,
            topics_json = ?,
            publication_year = ?,
            openalex_abstract = ?,
            updated_at = ?
        WHERE arxiv_id = ?
        """,
        (
            openalex_id,
            cited_by_count,
            fwci,
            _jdump(topics) if topics else None,
            publication_year,
            openalex_abstract,
            datetime.utcnow().isoformat(timespec="seconds"),
            arxiv_id,
        ),
    )


def update_methodology(
    conn: sqlite3.Connection,
    arxiv_id: str,
    tags: list[dict] | None,
) -> None:
    """
    Persiste las metodologias detectadas para un paper.
    tags = [{"family": str, "sem": float, "lex": int}, ...] ordenadas por confianza.
    No hace commit (el caller comitea por chunk).
    """
    conn.execute(
        "UPDATE papers SET methodology_json = ?, updated_at = ? WHERE arxiv_id = ?",
        (
            _jdump(tags) if tags else None,
            datetime.utcnow().isoformat(timespec="seconds"),
            arxiv_id,
        ),
    )


def all_methodology_families(conn: sqlite3.Connection) -> list[str]:
    """Lista deduplicada de familias presentes en el corpus (para poblar filtros)."""
    fams: set[str] = set()
    for (mj,) in conn.execute(
        "SELECT methodology_json FROM papers WHERE methodology_json IS NOT NULL"
    ):
        try:
            for t in json.loads(mj):
                f = t.get("family")
                if f:
                    fams.add(f)
        except Exception:
            pass
    return sorted(fams)


def update_artifacts(
    conn: sqlite3.Connection,
    arxiv_id: str,
    has_code: bool,
    has_data: bool,
    code_urls: list[str],
    data_urls: list[str],
) -> None:
    """Persiste el resultado del detector de artefactos."""
    conn.execute(
        """
        UPDATE papers SET
            has_code = ?,
            has_data = ?,
            code_urls_json = ?,
            data_urls_json = ?,
            updated_at = ?
        WHERE arxiv_id = ?
        """,
        (
            int(has_code),
            int(has_data),
            _jdump(code_urls) if code_urls else None,
            _jdump(data_urls) if data_urls else None,
            datetime.utcnow().isoformat(timespec="seconds"),
            arxiv_id,
        ),
    )


# ------------------------------------------------------------
# User ratings (likes / dislikes / notes)
# ------------------------------------------------------------

def rate_paper(
    conn: sqlite3.Connection,
    arxiv_id: str,
    rating: int,            # -1 dislike, 0 limpia, +1 like
    note: str | None = None,
) -> None:
    """Setea/limpia el rating del usuario. rating=0 limpia."""
    now = datetime.utcnow().isoformat(timespec="seconds")
    if rating == 0:
        conn.execute(
            "UPDATE papers SET user_rating = NULL, rated_at = ? WHERE arxiv_id = ?",
            (now, arxiv_id),
        )
    else:
        conn.execute(
            "UPDATE papers SET user_rating = ?, user_note = COALESCE(?, user_note), "
            "rated_at = ? WHERE arxiv_id = ?",
            (rating, note, now, arxiv_id),
        )
    conn.commit()


def set_note(conn: sqlite3.Connection, arxiv_id: str, note: str) -> None:
    """Guarda solo la nota libre, sin tocar el rating."""
    now = datetime.utcnow().isoformat(timespec="seconds")
    conn.execute(
        "UPDATE papers SET user_note = ?, rated_at = ? WHERE arxiv_id = ?",
        (note or None, now, arxiv_id),
    )
    conn.commit()


# ------------------------------------------------------------
# Embeddings
# ------------------------------------------------------------

def upsert_embedding(
    conn: sqlite3.Connection,
    arxiv_id: str,
    model: str,
    dim: int,
    vector_bytes: bytes,
) -> None:
    """Persiste un embedding ya normalizado L2 como BLOB float32."""
    conn.execute(
        """
        INSERT INTO embeddings (arxiv_id, model, dim, vector, computed_at)
        VALUES (?, ?, ?, ?, ?)
        ON CONFLICT(arxiv_id) DO UPDATE SET
            model = excluded.model,
            dim = excluded.dim,
            vector = excluded.vector,
            computed_at = excluded.computed_at
        """,
        (
            arxiv_id, model, dim, vector_bytes,
            datetime.utcnow().isoformat(timespec="seconds"),
        ),
    )


def papers_missing_embedding(
    conn: sqlite3.Connection,
    model: str,
) -> list[sqlite3.Row]:
    """Papers que aun no tienen embedding para ese modelo."""
    return conn.execute(
        """
        SELECT p.arxiv_id, p.title, p.abstract, p.openalex_abstract
        FROM papers p
        WHERE p.arxiv_id NOT IN (
            SELECT arxiv_id FROM embeddings WHERE model = ?
        )
        AND (p.title IS NOT NULL OR p.abstract IS NOT NULL)
        ORDER BY p.published DESC
        """,
        (model,),
    ).fetchall()


def load_all_embeddings(
    conn: sqlite3.Connection,
    model: str,
) -> list[sqlite3.Row]:
    """Carga todos los embeddings de un modelo (arxiv_id, vector blob)."""
    return conn.execute(
        "SELECT arxiv_id, vector FROM embeddings WHERE model = ?",
        (model,),
    ).fetchall()


def count_embeddings(conn: sqlite3.Connection, model: str | None = None) -> int:
    if model:
        return conn.execute(
            "SELECT COUNT(*) FROM embeddings WHERE model = ?", (model,)
        ).fetchone()[0]
    return conn.execute("SELECT COUNT(*) FROM embeddings").fetchone()[0]


def reset_openalex(conn: sqlite3.Connection) -> int:
    """
    Limpia todos los campos OpenAlex para forzar re-enrichment
    (util cuando se agregan campos nuevos al pipeline).
    Devuelve numero de papers reseteados.
    """
    cur = conn.execute(
        """
        UPDATE papers SET
            openalex_id = NULL,
            cited_by_count = NULL,
            fwci = NULL,
            topics_json = NULL,
            publication_year = NULL,
            openalex_abstract = NULL
        """
    )
    conn.commit()
    return cur.rowcount


def mark_openalex_unmatched(conn: sqlite3.Connection, arxiv_id: str) -> None:
    """
    Marca explicitamente un paper como buscado-pero-no-encontrado en OpenAlex.
    Usa openalex_id = 'NOT_FOUND' como sentinela para evitar re-querying en
    futuras corridas (no son free).
    """
    conn.execute(
        "UPDATE papers SET openalex_id = 'NOT_FOUND', updated_at = ? "
        "WHERE arxiv_id = ?",
        (datetime.utcnow().isoformat(timespec="seconds"), arxiv_id),
    )