"""
创建股票基础数据相关表
"""
from sqlalchemy import text
from app.db.session import SessionLocal

db = SessionLocal()

try:
    # 创建股票基础数据表
    db.execute(text("""
        CREATE TABLE IF NOT EXISTS stock_basic (
            id BIGSERIAL PRIMARY KEY,
            code VARCHAR(20) UNIQUE NOT NULL,
            name VARCHAR(50),
            total_shares BIGINT,
            float_shares BIGINT,
            industry_index_name VARCHAR(100),
            industry_index_code VARCHAR(20),
            institution_hold_ratio DECIMAL(10, 4),
            industry_level3 VARCHAR(100),
            list_date DATE,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )
    """))
    
    # 创建指数基础表
    db.execute(text("""
        CREATE TABLE IF NOT EXISTS index_basic (
            id BIGSERIAL PRIMARY KEY,
            code VARCHAR(20) UNIQUE NOT NULL,
            name VARCHAR(100),
            component_count INTEGER,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )
    """))
    
    # 创建指数交易表
    db.execute(text("""
        CREATE TABLE IF NOT EXISTS index_trade (
            id BIGSERIAL PRIMARY KEY,
            index_code VARCHAR(20) NOT NULL,
            trade_date DATE NOT NULL,
            open DECIMAL(10, 3),
            close DECIMAL(10, 3),
            high DECIMAL(10, 3),
            low DECIMAL(10, 3),
            change_pct DECIMAL(10, 4),
            volume BIGINT,
            amount DECIMAL(18, 2),
            total_market_value DECIMAL(18, 2),
            float_market_value DECIMAL(18, 2),
            up_count INTEGER,
            down_count INTEGER,
            flat_count INTEGER,
            limit_up_count INTEGER,
            limit_down_count INTEGER,
            suspend_count INTEGER,
            pe_ratio DECIMAL(10, 4),
            pe_median DECIMAL(10, 4),
            is_new_high BOOLEAN DEFAULT FALSE,
            is_new_low BOOLEAN DEFAULT FALSE,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            UNIQUE(index_code, trade_date)
        )
    """))
    
    # 创建索引
    db.execute(text("CREATE INDEX IF NOT EXISTS idx_stock_basic_code ON stock_basic(code)"))
    db.execute(text("CREATE INDEX IF NOT EXISTS idx_index_basic_code ON index_basic(code)"))
    db.execute(text("CREATE INDEX IF NOT EXISTS idx_index_trade_code ON index_trade(index_code)"))
    db.execute(text("CREATE INDEX IF NOT EXISTS idx_index_trade_date ON index_trade(trade_date)"))
    
    # 添加外键约束
    db.execute(text("""
        ALTER TABLE stock_basic 
        ADD CONSTRAINT fk_stock_basic_index_code 
        FOREIGN KEY (industry_index_code) REFERENCES index_basic(code)
    """))
    
    db.execute(text("""
        ALTER TABLE index_trade 
        ADD CONSTRAINT fk_index_trade_index_code 
        FOREIGN KEY (index_code) REFERENCES index_basic(code)
    """))
    
    db.commit()
    print("表创建成功")
except Exception as e:
    print(f"创建表失败: {str(e)}")
    db.rollback()
finally:
    db.close()