import os
import pypinyin
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse
import akshare as ak
import pandas as pd

app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

_stock_cache = None

def load_stock_cache():
    global _stock_cache
    if _stock_cache is None:
        try:
            df = ak.stock_info_a_code_name()
            _stock_cache = []
            for _, row in df.iterrows():
                code = str(row['code'])
                name = str(row['name'])
                py = "".join(pypinyin.lazy_pinyin(name, style=pypinyin.Style.FIRST_LETTER)).lower()
                _stock_cache.append({
                    "code": code,
                    "name": name,
                    "pinyin": py
                })
        except Exception as e:
            print("Error loading stock cache:", e)

def get_stock_name(symbol: str) -> str:
    load_stock_cache()
    if _stock_cache is not None:
        symbol = str(symbol)
        for item in _stock_cache:
            if item["code"] == symbol:
                return item["name"]
    return ""

def get_prefix_symbol(symbol: str) -> str:
    if len(symbol) == 8:
        return symbol
    if symbol.startswith('6'):
        return 'sh' + symbol
    return 'sz' + symbol

def format_stock_data(df: pd.DataFrame) -> list:
    formatted_data = []
    if df is None or df.empty:
        return formatted_data
        
    for _, row in df.iterrows():
        date_val = row.get('date', '')
        if hasattr(date_val, 'strftime'):
            date_str = date_val.strftime('%Y-%m-%d')
        else:
            date_str = str(date_val).split(' ')[0]
            
        formatted_data.append({
            'date': date_str,
            'open': float(row['open']),
            'close': float(row['close']),
            'low': float(row['low']),
            'high': float(row['high'])
        })
    return formatted_data

@app.get("/")
def serve_index():
    # 根路由直接返回 index.html
    html_path = os.path.join(os.path.dirname(__file__), "index.html")
    return FileResponse(html_path)

@app.get("/api/health")
def health_check():
    return {"status": "ok"}

@app.get("/api/stock")
def get_stock(symbol: str):
    full_symbol = get_prefix_symbol(symbol)
    stock_name = get_stock_name(symbol)
    try:
        df = ak.stock_zh_a_daily(symbol=full_symbol)
        data = format_stock_data(df)
        return {"name": stock_name, "data": data}
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"网络请求失败。原因: {str(e)}")

@app.get("/api/search")
def search_stock(q: str):
    load_stock_cache()
    if not _stock_cache:
        return []
    
    q_lower = q.lower()
    results = []
    for item in _stock_cache:
        if q_lower in item["code"] or q_lower in item["name"].lower() or q_lower in item["pinyin"]:
            results.append(item)
            if len(results) >= 10:
                break
    return results