from __future__ import annotations

import sqlite3

from backend.database.migrations.runner import Migration


def upgrade(connection: sqlite3.Connection) -> None:
    statements = (
        """
        CREATE TABLE screener_factor_snapshots (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            trade_date TEXT NOT NULL,
            version TEXT NOT NULL,
            observed_at TEXT NOT NULL,
            state TEXT NOT NULL CHECK (state IN ('final', 'archive')),
            source_set_json TEXT NOT NULL,
            coverage_json TEXT NOT NULL,
            created_at TEXT NOT NULL,
            UNIQUE (trade_date, version)
        )
        """,
        """
        CREATE TABLE screener_factor_values (
            snapshot_id INTEGER NOT NULL
                REFERENCES screener_factor_snapshots(id) ON DELETE CASCADE,
            identifier TEXT NOT NULL,
            code TEXT NOT NULL,
            name TEXT NOT NULL,
            sector TEXT,
            listed_days INTEGER NOT NULL,
            is_st INTEGER NOT NULL CHECK (is_st IN (0, 1)),
            payload_json TEXT NOT NULL,
            PRIMARY KEY (snapshot_id, identifier)
        )
        """,
        """
        CREATE TABLE screener_runs (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            owner_user_id INTEGER REFERENCES users(id) ON DELETE CASCADE,
            mode TEXT NOT NULL CHECK (mode IN ('stage', 'curated', 'custom')),
            strategy_id TEXT NOT NULL,
            strategy_name TEXT NOT NULL,
            strategy_version INTEGER NOT NULL,
            selection_date TEXT NOT NULL,
            factor_snapshot_id INTEGER NOT NULL
                REFERENCES screener_factor_snapshots(id) ON DELETE RESTRICT,
            status TEXT NOT NULL CHECK (
                status IN (
                    'pending', 'running', 'completed', 'no_signal',
                    'data_incomplete', 'failed'
                )
            ),
            started_at TEXT,
            completed_at TEXT,
            coverage REAL NOT NULL DEFAULT 0 CHECK (coverage >= 0 AND coverage <= 1),
            missing_fields_json TEXT NOT NULL DEFAULT '[]',
            result_json TEXT NOT NULL DEFAULT '[]',
            error_message TEXT NOT NULL DEFAULT ''
        )
        """,
        """
        CREATE UNIQUE INDEX screener_runs_idempotency_idx ON screener_runs (
            mode, strategy_id, selection_date, strategy_version,
            factor_snapshot_id, COALESCE(owner_user_id, 0)
        )
        """,
        """
        CREATE TABLE custom_screener_strategies (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            user_id INTEGER NOT NULL REFERENCES users(id) ON DELETE CASCADE,
            name TEXT NOT NULL,
            version INTEGER NOT NULL DEFAULT 1,
            formula_json TEXT NOT NULL,
            created_at TEXT NOT NULL,
            updated_at TEXT NOT NULL,
            UNIQUE (user_id, name)
        )
        """,
        """
        CREATE TABLE strategy_tracks (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            user_id INTEGER NOT NULL REFERENCES users(id) ON DELETE CASCADE,
            run_id INTEGER NOT NULL REFERENCES screener_runs(id) ON DELETE CASCADE,
            identifier TEXT NOT NULL,
            code TEXT NOT NULL,
            name TEXT NOT NULL,
            sector TEXT,
            selection_date TEXT NOT NULL,
            strategy_name TEXT NOT NULL,
            entry_price REAL NOT NULL CHECK (entry_price > 0),
            added_at TEXT NOT NULL,
            UNIQUE (user_id, run_id, identifier)
        )
        """,
        """
        CREATE TABLE strategy_track_bars (
            track_id INTEGER NOT NULL REFERENCES strategy_tracks(id) ON DELETE CASCADE,
            trade_date TEXT NOT NULL,
            open REAL NOT NULL,
            high REAL NOT NULL,
            low REAL NOT NULL,
            close REAL NOT NULL,
            PRIMARY KEY (track_id, trade_date)
        )
        """,
        """
        CREATE TABLE strategy_track_events (
            track_id INTEGER NOT NULL REFERENCES strategy_tracks(id) ON DELETE CASCADE,
            milestone TEXT NOT NULL CHECK (milestone IN ('t1', 't5')),
            created_at TEXT NOT NULL,
            PRIMARY KEY (track_id, milestone)
        )
        """,
    )
    for statement in statements:
        connection.execute(statement)


def downgrade(connection: sqlite3.Connection) -> None:
    for table in (
        "strategy_track_events",
        "strategy_track_bars",
        "strategy_tracks",
        "custom_screener_strategies",
        "screener_runs",
        "screener_factor_values",
        "screener_factor_snapshots",
    ):
        connection.execute(f"DROP TABLE {table}")


MIGRATION = Migration(
    version=7,
    name="create_deterministic_screener",
    signature="screener:v1:factors-runs-custom-tracking",
    upgrade=upgrade,
    downgrade=downgrade,
)