import asyncio
import aiosqlite
import json
import time
from pathlib import Path
from typing import Optional
from enum import Enum

import structlog

from src.config import settings

logger = structlog.get_logger()


class ResearchStatus(str, Enum):
    RUNNING = "running"
    SATURATED = "saturated"
    FINISHED = "finished"
    ERROR = "error"


class OutputType(str, Enum):
    PODCAST = "podcast"
    BLOG = "blog"
    REPORT = "report"
    THREAD = "thread"
    REPORT_EXTENDED = "report_extended"
    BLOG_EXTENDED = "blog_extended"
    PODCAST_EXTENDED = "podcast_extended"


SCHEMA = """
CREATE TABLE IF NOT EXISTS research_sessions (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    topic TEXT NOT NULL,
    status TEXT NOT NULL DEFAULT 'running',
    telegram_chat_id INTEGER NOT NULL,
    telegram_message_id INTEGER,
    created_at REAL NOT NULL,
    updated_at REAL NOT NULL,
    iterations INTEGER DEFAULT 0,
    total_sources INTEGER DEFAULT 0,
    total_chunks INTEGER DEFAULT 0,
    total_words INTEGER DEFAULT 0,
    meta JSON DEFAULT '{}'
);

CREATE TABLE IF NOT EXISTS sources (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    session_id INTEGER NOT NULL REFERENCES research_sessions(id),
    url TEXT NOT NULL,
    title TEXT,
    source_type TEXT,          -- wikipedia, reddit, youtube, pdf, web, rss
    depth INTEGER DEFAULT 0,
    quality_score REAL DEFAULT 0,
    word_count INTEGER DEFAULT 0,
    scraped_at REAL,
    status TEXT DEFAULT 'pending',  -- pending, scraped, failed, skipped
    error TEXT,
    UNIQUE(session_id, url)
);

CREATE TABLE IF NOT EXISTS chunks (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    session_id INTEGER NOT NULL REFERENCES research_sessions(id),
    source_id INTEGER NOT NULL REFERENCES sources(id),
    content TEXT NOT NULL,
    chunk_index INTEGER NOT NULL,
    token_count INTEGER,
    quality_score REAL DEFAULT 0,
    embedding JSON,            -- stored as JSON array for sqlite-vec compat
    created_at REAL NOT NULL
);

CREATE TABLE IF NOT EXISTS outputs (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    session_id INTEGER NOT NULL REFERENCES research_sessions(id),
    output_type TEXT NOT NULL,
    content TEXT NOT NULL,
    created_at REAL NOT NULL
);

CREATE TABLE IF NOT EXISTS source_contents (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    source_id INTEGER NOT NULL UNIQUE REFERENCES sources(id),
    content TEXT NOT NULL,
    created_at REAL NOT NULL
);

CREATE INDEX IF NOT EXISTS idx_sources_session ON sources(session_id);
CREATE INDEX IF NOT EXISTS idx_chunks_session ON chunks(session_id);
CREATE INDEX IF NOT EXISTS idx_chunks_quality ON chunks(session_id, quality_score DESC);
CREATE INDEX IF NOT EXISTS idx_source_contents ON source_contents(source_id);

CREATE TABLE IF NOT EXISTS api_usage (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    session_id INTEGER REFERENCES research_sessions(id),
    call_type TEXT NOT NULL,
    model TEXT NOT NULL,
    input_tokens INTEGER NOT NULL,
    output_tokens INTEGER NOT NULL,
    cost_usd REAL NOT NULL,
    created_at REAL NOT NULL
);

CREATE TABLE IF NOT EXISTS watched_topics (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    topic TEXT NOT NULL,
    chat_id INTEGER NOT NULL,
    interval_hours INTEGER NOT NULL DEFAULT 24,
    next_run_at REAL NOT NULL,
    last_run_at REAL,
    enabled INTEGER NOT NULL DEFAULT 1,
    created_at REAL NOT NULL,
    UNIQUE(topic, chat_id)
);

CREATE TABLE IF NOT EXISTS news_seen (
    id INTEGER PRIMARY KEY,
    source TEXT NOT NULL,
    guid TEXT NOT NULL,
    title TEXT,
    link TEXT,
    published_at TEXT,
    matched_keywords TEXT,
    seen_at TEXT DEFAULT (datetime('now')),
    notified INTEGER DEFAULT 0,
    UNIQUE(source, guid)
);

CREATE INDEX IF NOT EXISTS idx_news_seen_source ON news_seen(source);
CREATE INDEX IF NOT EXISTS idx_news_seen_seen ON news_seen(seen_at);
"""


# Conexión única compartida por proceso. aiosqlite serializa todas las
# operaciones en el hilo de la conexión, así que compartirla entre coroutines
# es seguro y además evita la contención del lock de escritura a nivel de
# fichero que sufrían las conexiones múltiples (scheduler, /compare).
_shared_conn: Optional[aiosqlite.Connection] = None
_init_lock = asyncio.Lock()


class _SharedConnection:
    """Proxy sobre la conexión compartida cuyo close() es no-op.

    Permite que los handlers sigan usando el patrón `db = await get_db()` …
    `finally: await db.close()` sin cambios, mientras por debajo se reutiliza
    una sola conexión real durante toda la vida del proceso.
    """
    def __init__(self, conn: aiosqlite.Connection):
        self._conn = conn

    def __getattr__(self, name):
        return getattr(self._conn, name)

    async def close(self):
        # No cerrar: la conexión es compartida por todo el proceso.
        pass


async def _init_shared() -> aiosqlite.Connection:
    global _shared_conn
    if _shared_conn is None:
        async with _init_lock:
            if _shared_conn is None:  # doble check tras adquirir el lock
                Path(settings.db_path).parent.mkdir(parents=True, exist_ok=True)
                conn = await aiosqlite.connect(settings.db_path)
                conn.row_factory = aiosqlite.Row
                await conn.execute("PRAGMA journal_mode=WAL")
                await conn.execute("PRAGMA synchronous=NORMAL")
                # Espera hasta 5s si OTRO proceso tiene el lock de escritura
                # antes de fallar con "database is locked".
                await conn.execute("PRAGMA busy_timeout=5000")
                await conn.executescript(SCHEMA)
                await conn.commit()
                _shared_conn = conn
                logger.info("Shared DB connection initialized", path=settings.db_path)
    return _shared_conn


async def get_db() -> aiosqlite.Connection:
    conn = await _init_shared()
    return _SharedConnection(conn)


async def close_db() -> None:
    """Cierra la conexión compartida (checkpoint WAL). Llamar al apagar el bot."""
    global _shared_conn
    if _shared_conn is not None:
        try:
            await _shared_conn.close()
        finally:
            _shared_conn = None


class ResearchDB:
    def __init__(self, db: aiosqlite.Connection):
        self.db = db

    # --- Sessions ---

    async def create_session(self, topic: str, chat_id: int) -> int:
        now = time.time()
        cursor = await self.db.execute(
            """INSERT INTO research_sessions (topic, status, telegram_chat_id, created_at, updated_at)
               VALUES (?, ?, ?, ?, ?)""",
            (topic, ResearchStatus.RUNNING, chat_id, now, now)
        )
        await self.db.commit()
        return cursor.lastrowid

    async def get_session(self, session_id: int) -> Optional[dict]:
        cursor = await self.db.execute(
            "SELECT * FROM research_sessions WHERE id = ?", (session_id,)
        )
        row = await cursor.fetchone()
        return dict(row) if row else None

    async def get_latest_session(self, chat_id: int) -> Optional[dict]:
        cursor = await self.db.execute(
            "SELECT * FROM research_sessions WHERE telegram_chat_id = ? ORDER BY created_at DESC LIMIT 1",
            (chat_id,)
        )
        row = await cursor.fetchone()
        return dict(row) if row else None

    async def get_session_urls(self, session_id: int) -> set:
        async with self.db.execute(
            "SELECT url FROM sources WHERE session_id = ?", (session_id,)
        ) as cur:
            rows = await cur.fetchall()
            return {r[0] for r in rows}

    async def get_previous_session(self, chat_id: int, topic: str,
                                    exclude_session_id: int) -> Optional[dict]:
        async with self.db.execute(
            """SELECT id, topic, status, created_at FROM research_sessions
               WHERE telegram_chat_id = ? AND topic = ? AND id != ?
               ORDER BY created_at DESC LIMIT 1""",
            (chat_id, topic, exclude_session_id)
        ) as cur:
            row = await cur.fetchone()
            if not row:
                return None
            return {"id": row[0], "topic": row[1],
                    "status": row[2], "created_at": row[3]}

    async def get_active_session(self, chat_id: int) -> Optional[dict]:
        cursor = await self.db.execute(
            """SELECT * FROM research_sessions 
               WHERE telegram_chat_id = ? AND status = 'running'
               ORDER BY created_at DESC LIMIT 1""",
            (chat_id,)
        )
        row = await cursor.fetchone()
        return dict(row) if row else None

    async def update_session(self, session_id: int, **kwargs):
        kwargs["updated_at"] = time.time()
        sets = ", ".join(f"{k} = ?" for k in kwargs)
        values = list(kwargs.values()) + [session_id]
        await self.db.execute(
            f"UPDATE research_sessions SET {sets} WHERE id = ?", values
        )
        await self.db.commit()

    async def get_session_stats(self, session_id: int) -> dict:
        cursor = await self.db.execute(
            """SELECT
                COUNT(*) as total,
                SUM(CASE WHEN status='scraped'  THEN 1 ELSE 0 END) as scraped,
                SUM(CASE WHEN status='failed'   THEN 1 ELSE 0 END) as failed,
                SUM(CASE WHEN status='pending'  THEN 1 ELSE 0 END) as pending,
                SUM(CASE WHEN status='skipped'  THEN 1 ELSE 0 END) as skipped
               FROM sources WHERE session_id = ?""",
            (session_id,)
        )
        row = await cursor.fetchone()
        return dict(row) if row else {}

    # --- Sources ---

    async def add_source(self, session_id: int, url: str, source_type: str,
                         depth: int = 0, title: str = None) -> Optional[int]:
        try:
            cursor = await self.db.execute(
                """INSERT OR IGNORE INTO sources (session_id, url, title, source_type, depth)
                   VALUES (?, ?, ?, ?, ?)""",
                (session_id, url, title, source_type, depth)
            )
            await self.db.commit()
            return cursor.lastrowid if cursor.rowcount > 0 else None
        except Exception:
            return None

    async def update_source(self, source_id: int, **kwargs):
        sets = ", ".join(f"{k} = ?" for k in kwargs)
        values = list(kwargs.values()) + [source_id]
        await self.db.execute(f"UPDATE sources SET {sets} WHERE id = ?", values)
        await self.db.commit()

    async def get_pending_sources(self, session_id: int, limit: int = 10) -> list[dict]:
        cursor = await self.db.execute(
            """SELECT * FROM sources WHERE session_id = ? AND status = 'pending'
               ORDER BY depth ASC, id ASC LIMIT ?""",
            (session_id, limit)
        )
        rows = await cursor.fetchall()
        return [dict(r) for r in rows]

    async def get_all_sources(self, session_id: int) -> list[dict]:
        cursor = await self.db.execute(
            "SELECT * FROM sources WHERE session_id = ? ORDER BY quality_score DESC",
            (session_id,)
        )
        rows = await cursor.fetchall()
        return [dict(r) for r in rows]

    async def source_exists(self, session_id: int, url: str) -> bool:
        cursor = await self.db.execute(
            "SELECT 1 FROM sources WHERE session_id = ? AND url = ?",
            (session_id, url)
        )
        return await cursor.fetchone() is not None

    # --- Chunks ---

    async def add_chunk(self, session_id: int, source_id: int, content: str,
                        chunk_index: int, token_count: int, quality_score: float,
                        embedding: Optional[list] = None) -> int:
        cursor = await self.db.execute(
            """INSERT INTO chunks (session_id, source_id, content, chunk_index,
               token_count, quality_score, embedding, created_at)
               VALUES (?, ?, ?, ?, ?, ?, ?, ?)""",
            (session_id, source_id, content, chunk_index,
             token_count, quality_score,
             json.dumps(embedding) if embedding else None,
             time.time())
        )
        await self.db.commit()
        return cursor.lastrowid

    async def get_top_chunks(self, session_id: int, limit: int = 50) -> list[dict]:
        cursor = await self.db.execute(
            """SELECT c.*, s.url, s.title, s.source_type FROM chunks c
               JOIN sources s ON c.source_id = s.id
               WHERE c.session_id = ? AND c.quality_score >= ?
               ORDER BY c.quality_score DESC LIMIT ?""",
            (session_id, settings.quality_threshold, limit)
        )
        rows = await cursor.fetchall()
        return [dict(r) for r in rows]

    async def get_chunks_count(self, session_id: int) -> int:
        cursor = await self.db.execute(
            "SELECT COUNT(*) FROM chunks WHERE session_id = ?", (session_id,)
        )
        row = await cursor.fetchone()
        return row[0]

    # --- Outputs ---

    async def save_output(self, session_id: int, output_type: str, content: str) -> int:
        cursor = await self.db.execute(
            "INSERT INTO outputs (session_id, output_type, content, created_at) VALUES (?, ?, ?, ?)",
            (session_id, output_type, content, time.time())
        )
        await self.db.commit()
        return cursor.lastrowid

    async def save_source_content(self, source_id: int, content: str):
        await self.db.execute(
            """INSERT OR REPLACE INTO source_contents (source_id, content, created_at)
               VALUES (?, ?, ?)""",
            (source_id, content, time.time())
        )
        await self.db.commit()

    async def get_source_content(self, source_id: int) -> Optional[str]:
        cursor = await self.db.execute(
            "SELECT content FROM source_contents WHERE source_id = ?", (source_id,)
        )
        row = await cursor.fetchone()
        return row[0] if row else None

    async def get_cached_content(self, url: str,
                                  max_age_days: int = 7) -> Optional[str]:
        threshold = time.time() - (max_age_days * 86400)
        async with self.db.execute(
            """SELECT sc.content FROM source_contents sc
               JOIN sources s ON s.id = sc.source_id
               WHERE s.url = ? AND sc.created_at > ?
               ORDER BY sc.created_at DESC LIMIT 1""",
            (url, threshold)
        ) as cur:
            row = await cur.fetchone()
            return row[0] if row else None

    async def get_outputs(self, session_id: int) -> list[dict]:
        cursor = await self.db.execute(
            "SELECT * FROM outputs WHERE session_id = ? ORDER BY created_at DESC",
            (session_id,)
        )
        rows = await cursor.fetchall()
        return [dict(r) for r in rows]

    # --- API Usage ---

    # Precios Claude en USD por 1M de tokens (input, output).
    # Se busca por substring del id del modelo; fallback a Haiku.
    _MODEL_PRICING = {
        "opus":   (15.00, 75.00),
        "sonnet": (3.00, 15.00),
        "haiku":  (0.80, 4.00),
    }

    @classmethod
    def _price_for_model(cls, model: str) -> tuple[float, float]:
        m = (model or "").lower()
        for key, price in cls._MODEL_PRICING.items():
            if key in m:
                return price
        return cls._MODEL_PRICING["haiku"]

    async def log_api_call(self, session_id, call_type: str, model: str,
                            input_tokens: int, output_tokens: int):
        in_price, out_price = self._price_for_model(model)
        cost = (input_tokens * in_price + output_tokens * out_price) / 1_000_000
        await self.db.execute(
            """INSERT INTO api_usage
               (session_id, call_type, model, input_tokens, output_tokens, cost_usd, created_at)
               VALUES (?,?,?,?,?,?,?)""",
            (session_id, call_type, model, input_tokens, output_tokens, cost, time.time())
        )
        await self.db.commit()

    async def get_usage_stats(self, session_id: int) -> list[dict]:
        cursor = await self.db.execute(
            """SELECT call_type,
                      COUNT(*) as calls,
                      SUM(input_tokens + output_tokens) as total_tokens,
                      SUM(cost_usd) as total_cost
               FROM api_usage WHERE session_id = ?
               GROUP BY call_type""",
            (session_id,)
        )
        rows = await cursor.fetchall()
        return [dict(r) for r in rows]

    async def get_total_usage_stats(self) -> dict:
        cursor = await self.db.execute(
            """SELECT COUNT(DISTINCT session_id) as sessions,
                      SUM(cost_usd) as total_cost
               FROM api_usage"""
        )
        row = await cursor.fetchone()
        return dict(row) if row else {"sessions": 0, "total_cost": 0}

    # --- Watched Topics ---

    async def add_watch(self, topic: str, chat_id: int, interval_hours: int,
                        next_run_at: Optional[float] = None) -> int:
        now = time.time()
        if next_run_at is None:
            next_run_at = now + interval_hours * 3600
        cursor = await self.db.execute(
            """INSERT OR REPLACE INTO watched_topics
               (topic, chat_id, interval_hours, next_run_at, created_at)
               VALUES (?, ?, ?, ?, ?)""",
            (topic, chat_id, interval_hours, next_run_at, now)
        )
        await self.db.commit()
        return cursor.lastrowid

    async def remove_watch(self, topic: str, chat_id: int) -> bool:
        cursor = await self.db.execute(
            "DELETE FROM watched_topics WHERE topic = ? AND chat_id = ?",
            (topic, chat_id)
        )
        await self.db.commit()
        return cursor.rowcount > 0

    async def list_watches(self, chat_id: int) -> list[dict]:
        cursor = await self.db.execute(
            "SELECT * FROM watched_topics WHERE chat_id = ? ORDER BY created_at ASC",
            (chat_id,)
        )
        rows = await cursor.fetchall()
        return [dict(r) for r in rows]

    async def get_due_watches(self) -> list[dict]:
        cursor = await self.db.execute(
            "SELECT * FROM watched_topics WHERE enabled = 1 AND next_run_at <= ?",
            (time.time(),)
        )
        rows = await cursor.fetchall()
        return [dict(r) for r in rows]

    async def update_watch_run(self, watch_id: int):
        cursor = await self.db.execute(
            "SELECT interval_hours FROM watched_topics WHERE id = ?", (watch_id,)
        )
        row = await cursor.fetchone()
        if not row:
            return
        now = time.time()
        await self.db.execute(
            "UPDATE watched_topics SET last_run_at = ?, next_run_at = ? WHERE id = ?",
            (now, now + row[0] * 3600, watch_id)
        )
        await self.db.commit()

    # --- News monitor ---

    async def source_seeded(self, source: str) -> bool:
        """True si ya hemos visto algún item de esta fuente (cold-start por fuente:
        el primer poll siembra sin notificar)."""
        cursor = await self.db.execute(
            "SELECT 1 FROM news_seen WHERE source = ? LIMIT 1", (source,)
        )
        return await cursor.fetchone() is not None

    async def record_news_item(self, source: str, guid: str, title: str,
                               link: str, published_at: str, matched: str,
                               notified: int) -> bool:
        """INSERT OR IGNORE de un item. Devuelve True solo si es una inserción
        real (item nuevo); False si ya existía (UNIQUE(source, guid))."""
        cursor = await self.db.execute(
            """INSERT OR IGNORE INTO news_seen
               (source, guid, title, link, published_at, matched_keywords, notified)
               VALUES (?, ?, ?, ?, ?, ?, ?)""",
            (source, guid, title, link, published_at, matched, notified)
        )
        await self.db.commit()
        return cursor.rowcount > 0

    async def mark_news_notified(self, ids: list[int]):
        if not ids:
            return
        placeholders = ",".join("?" for _ in ids)
        await self.db.execute(
            f"UPDATE news_seen SET notified = 1 WHERE id IN ({placeholders})",
            tuple(ids)
        )
        await self.db.commit()

    async def get_recent_news(self, hours: int = 24) -> list[dict]:
        """Items matcheados cuya seen_at (o published_at parseable) cae en la
        ventana. Orden por published_at desc, con seen_at como desempate."""
        window = f"-{int(hours)} hours"
        cursor = await self.db.execute(
            """SELECT * FROM news_seen
               WHERE seen_at >= datetime('now', ?)
                  OR (published_at IS NOT NULL AND published_at != ''
                      AND datetime(published_at) >= datetime('now', ?))
               ORDER BY published_at DESC, seen_at DESC
               LIMIT 50""",
            (window, window)
        )
        rows = await cursor.fetchall()
        return [dict(r) for r in rows]

    async def get_unnotified(self, limit: Optional[int] = None) -> list[dict]:
        """Items aún sin notificar (notified=0), más recientes primero. El
        scheduler (F2) los empuja a Telegram y luego los marca con
        mark_news_notified. Aplica LIMIT solo si `limit` no es None."""
        sql = (
            "SELECT id, source, title, link, published_at, matched_keywords "
            "FROM news_seen WHERE notified = 0 "
            "ORDER BY published_at DESC NULLS LAST, seen_at DESC"
        )
        params: tuple = ()
        if limit is not None:
            sql += " LIMIT ?"
            params = (int(limit),)
        cursor = await self.db.execute(sql, params)
        rows = await cursor.fetchall()
        return [dict(r) for r in rows]

    # --- Maintenance ---

    async def purge_old_sessions(self, max_age_days: int = 30) -> dict:
        await self.db.execute("PRAGMA foreign_keys = ON")

        threshold = time.time() - max_age_days * 86400
        cursor = await self.db.execute(
            "SELECT id FROM research_sessions WHERE created_at < ? AND status != 'running'",
            (threshold,)
        )
        session_ids = [row[0] for row in await cursor.fetchall()]

        counts = {"sessions": 0, "sources": 0, "chunks": 0, "outputs": 0, "api_usage": 0}

        for sid in session_ids:
            await self.db.execute(
                "DELETE FROM source_contents WHERE source_id IN (SELECT id FROM sources WHERE session_id = ?)",
                (sid,)
            )
            cur = await self.db.execute("DELETE FROM chunks WHERE session_id = ?", (sid,))
            counts["chunks"] += cur.rowcount
            cur = await self.db.execute("DELETE FROM outputs WHERE session_id = ?", (sid,))
            counts["outputs"] += cur.rowcount
            cur = await self.db.execute("DELETE FROM api_usage WHERE session_id = ?", (sid,))
            counts["api_usage"] += cur.rowcount
            cur = await self.db.execute("DELETE FROM sources WHERE session_id = ?", (sid,))
            counts["sources"] += cur.rowcount
            cur = await self.db.execute("DELETE FROM research_sessions WHERE id = ?", (sid,))
            counts["sessions"] += cur.rowcount

        await self.db.commit()
        logger.info("Purged sessions older than days",
                    sessions=counts["sessions"], days=max_age_days)
        return counts