import os
import sys
import time
import torch
import torch.nn as nn
import pandas as pd
import numpy as np
import json
import random
from loguru import logger
from datetime import datetime, timedelta
from sentence_transformers import SentenceTransformer
from dotenv import load_dotenv

load_dotenv(os.path.expanduser("~/.config/opencode/.env"))

# Setup paths
KRONOS_DIR = os.path.dirname(os.path.abspath(__file__))
SRC_DIR = os.path.dirname(os.path.dirname(KRONOS_DIR))
if SRC_DIR not in sys.path:
    sys.path.insert(0, SRC_DIR)

from ..kronos.model import Kronos, KronosTokenizer, KronosPredictor
from ..database_manager import DatabaseManager
from ..stock_tools import StockTools
from ..search_tools import SearchTools
from ..llm.factory import get_model
from ..visualizer import VisualizerTools
from ..schema.models import ForecastResult, KLinePoint
from agno.agent import Agent


class AutoSynthesisTrainer:
    def __init__(self, news_dim=384):
        self.device = (
            "cuda"
            if torch.cuda.is_available()
            else "mps"
            if torch.backends.mps.is_available()
            else "cpu"
        )
        self.db = DatabaseManager()
        self.tools = StockTools(self.db)
        self.searcher = SearchTools(self.db)
        # Try loading from local cache first to avoid network timeouts
        model_name = os.getenv(
            "EMBEDDING_MODEL", "sentence-transformers/all-MiniLM-L6-v2"
        )
        try:
            logger.info(f"🔄 Attempting to load {model_name} from local cache...")
            self.embedder = SentenceTransformer(
                model_name, device=self.device, local_files_only=True
            )
            logger.success("✅ Model loaded from local cache.")
        except Exception:
            logger.warning(
                "⚠️ Local cache not found or incomplete. Attempting to download..."
            )
            self.embedder = SentenceTransformer(model_name, device=self.device)
        self.news_dim = news_dim

        # Try loading from local cache first to avoid network timeouts
        try:
            logger.info(
                "🔄 Attempting to load Kronos and Tokenizer from local cache..."
            )
            self.tokenizer = KronosTokenizer.from_pretrained(
                "NeoQuasar/Kronos-Tokenizer-base", local_files_only=True
            ).to(self.device)
            base_model = Kronos.from_pretrained(
                "NeoQuasar/Kronos-base", local_files_only=True
            )
            logger.success("✅ Kronos and Tokenizer loaded from local cache.")
        except Exception:
            logger.warning(
                "⚠️ Local Kronos/Tokenizer not found or incomplete. Attempting to download..."
            )
            self.tokenizer = KronosTokenizer.from_pretrained(
                "NeoQuasar/Kronos-Tokenizer-base"
            ).to(self.device)
            base_model = Kronos.from_pretrained("NeoQuasar/Kronos-base")

        self.model = Kronos(
            base_model.s1_bits,
            base_model.s2_bits,
            base_model.n_layers,
            base_model.d_model,
            base_model.n_heads,
            base_model.ff_dim,
            base_model.ffn_dropout_p,
            base_model.attn_dropout_p,
            base_model.resid_dropout_p,
            base_model.token_dropout_p,
            base_model.learn_te,
            news_dim=self.news_dim,
        ).to(self.device)
        self.model.load_state_dict(base_model.state_dict(), strict=False)

        # LLM for causality verification
        provider = os.getenv("LLM_PROVIDER", "minimax")
        model_id = os.getenv("LLM_MODEL", "Qwen")
        self.llm_agent = Agent(model=get_model(provider, model_id))

    def discover_shocks(
        self, ticker_list, threshold=2.0, limit_per_stock=5, days=365, pred_len=5
    ):
        """1. Find days with significant price movements (Look back 1 year)"""
        shocks = []
        end_date = datetime.now().strftime("%Y-%m-%d")
        start_date = (datetime.now() - timedelta(days=days)).strftime("%Y-%m-%d")

        for ticker in ticker_list:
            df = self.tools.get_stock_price(
                ticker, start_date=start_date, end_date=end_date
            )
            if df.empty or len(df) < 60:
                continue

            # Look for big moves
            moves = df[df["change_pct"].abs() > threshold].copy()
            if moves.empty:
                continue

            count = 0
            for idx, row in moves.iterrows():
                # Ensure we have history before this day AND enough future days for eval
                date_idx = df.index.get_loc(idx)
                if date_idx < 50 or date_idx + pred_len > len(df):
                    continue

                shocks.append(
                    {
                        "ticker": ticker,
                        "date": row["date"],
                        "change": row["change_pct"],
                        "history": df.iloc[date_idx - 50 : date_idx],
                        "target": df.iloc[
                            date_idx : date_idx + pred_len
                        ],  # Now capturing pred_len days
                    }
                )
                count += 1
                if count >= limit_per_stock:
                    break

        logger.info(
            f"✨ Discovered {len(shocks)} potential price shocks over the last {days} days."
        )
        return shocks

    def find_reason_and_verify(self, shock):
        """2. Search for reasons and verify causality using LLM"""
        ticker_info = self.db.get_stock_by_code(shock["ticker"])
        name = ticker_info["name"] if ticker_info else shock["ticker"]
        date_str = shock["date"]

        # Try multiple query variations and engines
        queries = [
            f"{name} ({shock['ticker']}) {date_str} 为什么涨跌 原因",
            f"{name} {date_str} 异动 原因",
            f"{shock['ticker']} {date_str} 新闻",
        ]

        search_results = []
        for query in queries:
            logger.info(f"🔍 Searching for reason: {query}")
            # Try alternate engines
            for engine in ["baidu"]:
                try:
                    results = self.searcher.search_list(
                        query, engine=engine, max_results=3, enrich=False
                    )
                    if results:
                        search_results = results
                        break
                except Exception as e:
                    logger.warning(f"Search failed for {query} on {engine}: {e}")

            if search_results:
                break
            time.sleep(random.uniform(1.0, 2.0))

        if not search_results:
            logger.warning(
                f"⚠️ No search results found for {name} on {date_str} after multiple attempts."
            )
            return None

        context = "\n".join(
            [f"- {r['title']}: {r.get('content', '')[:300]}" for r in search_results]
        )

        prompt = f"""
        任务：判断以下新闻是否解释了该股票在 {date_str} 的 {shock["change"]:.2f}% 价格变动。
        
        股票：{name}
        日期：{date_str}
        变动：{shock["change"]:.2f}%
        
        搜索结果：
        {context}
        
        要求：
        1. 该新闻是否在该日期左右发生？
        2. 该新闻是否能逻辑上解释这种大幅波动（如财报、利好政策、重组、大环境暴跌等）？
        3. 如果是，请总结一段 100 字以内的“核心推动原因”。
        4. 返回 JSON: {{"is_causal": true/false, "summary": "原因摘要"}}
        """

        try:
            res = self.llm_agent.run(prompt)
            data = json.loads(
                res.content.replace("```json", "").replace("```", "").strip()
            )
            if data.get("is_causal"):
                logger.success(
                    f"✅ Verified cause for {name} on {date_str}: {data['summary']}"
                )
                return data["summary"]
            else:
                logger.warning(
                    f"❌ Verified cause for {name} on {date_str}: {data['summary']}"
                )
                return None
        except Exception as e:
            logger.warning(f"Verification failed: {e}")
        return None

    def save_model(self, path=None):
        """Save the news_proj weights"""
        if path is None:
            save_dir = os.path.join(SRC_DIR, "exports/models")
            os.makedirs(save_dir, exist_ok=True)
            path = os.path.join(
                save_dir, f"kronos_news_v1_{datetime.now().strftime('%Y%m%d_%H%M')}.pt"
            )

        # We only really need to save the news_proj part as it's the only one we train
        torch.save(
            {
                "news_proj_state_dict": self.model.news_proj.state_dict(),
                "news_dim": self.news_dim,
                "d_model": self.model.d_model,
            },
            path,
        )
        logger.success(f"💾 Model weights saved to {path}")
        return path

    def run_synthesis_and_train(self, tickers, pred_len=5):
        # 1. Discovery
        shocks = self.discover_shocks(tickers, pred_len=pred_len)
        print(f"find {len(shocks)} shocks")

        # 2. News Association & Verification
        dataset = []
        max_news_items = 200  # Limit to 200 news items per session to avoid search bans

        logger.info(
            f"🧬 Starting News Association for {len(shocks)} shocks (Max limit: {max_news_items})"
        )

        for i, shock in enumerate(shocks):
            if len(dataset) >= max_news_items:
                logger.info("Reached maximum news items limit for this session.")
                break

            summary = self.find_reason_and_verify(shock)
            if summary:
                # 3. Embedding news
                emb = self.embedder.encode(summary)
                dataset.append(
                    {
                        "history": shock["history"],
                        "target": shock["target"],
                        "news_emb": emb,
                        "summary": summary,
                    }
                )

            # Add delay after search with randomness to avoid being blocked
            if i < len(shocks) - 1:
                delay = random.uniform(2.0, 4.0)
                time.sleep(delay)

        if not dataset:
            logger.error(
                "❌ No verified news-price pairs found. Adjust threshold or check if news is available in that period."
            )
            return

        # 4. Train/Val Split
        random.seed(42)
        random.shuffle(dataset)

        if len(dataset) < 2:
            train_set = dataset
            val_set = []
            logger.warning(
                f"⚠️ Only {len(dataset)} sample(s) found. Training on all, skipping validation."
            )
        else:
            split_idx = max(1, int(len(dataset) * 0.8))
            if split_idx >= len(dataset):
                split_idx = len(dataset) - 1

            train_set = dataset[:split_idx]
            val_set = dataset[split_idx:]
            logger.info(
                f"🏗️ Dataset Split: {len(train_set)} samples for training, {len(val_set)} for validation."
            )

        if not train_set:
            logger.error("❌ No samples for training.")
            return

        # 5. Training (Few-shot)
        optimizer = torch.optim.Adam(self.model.news_proj.parameters(), lr=1e-3)
        criterion = nn.CrossEntropyLoss()
        self.model.train()

        loss_history = []
        logger.info(f"🚀 Training for 30 epochs...")
        for epoch in range(30):
            total_loss = 0
            for item in train_set:
                optimizer.zero_grad()

                # Prep Data
                hist_df = item["history"]
                # For training, we still focus on the immediate next point (teacher forcing)
                target_df = item["target"].iloc[:1]

                hist_raw = hist_df[
                    ["open", "high", "low", "close", "volume"]
                ].values.astype(np.float32)
                hist_raw = np.column_stack([hist_raw, hist_raw[:, 3] * hist_raw[:, 4]])

                mean, std = hist_raw.mean(axis=0), hist_raw.std(axis=0) + 1e-5
                hist_norm = (
                    torch.from_numpy((hist_raw - mean) / std)
                    .unsqueeze(0)
                    .to(self.device)
                )

                target_raw = target_df[
                    ["open", "high", "low", "close", "volume"]
                ].values.astype(np.float32)
                target_raw = np.column_stack(
                    [target_raw, target_raw[:, 3] * target_raw[:, 4]]
                )
                target_norm = (
                    torch.from_numpy((target_raw - mean) / std)
                    .unsqueeze(0)
                    .to(self.device)
                )

                with torch.no_grad():
                    z_indices = self.tokenizer.encode(hist_norm, half=True)
                    t_indices = self.tokenizer.encode(target_norm, half=True)
                    s1_ids, s2_ids = z_indices[0], z_indices[1]
                    t_s1, t_s2 = t_indices[0], t_indices[1]

                news_t = torch.from_numpy(item["news_emb"]).unsqueeze(0).to(self.device)
                s1_logits, s2_logits = self.model(
                    s1_ids,
                    s2_ids,
                    news_emb=news_t,
                    use_teacher_forcing=True,
                    s1_targets=t_s1,
                )

                loss = (
                    criterion(s1_logits[:, -1, :], t_s1[:, 0])
                    + criterion(s2_logits[:, -1, :], t_s2[:, 0])
                ) / 2
                loss.backward()
                optimizer.step()
                total_loss += loss.item()

            avg_epoch_loss = total_loss / max(1, len(train_set))
            loss_history.append(avg_epoch_loss)

            if (epoch + 1) % 10 == 0:
                logger.info(f"Epoch {epoch + 1} Loss: {avg_epoch_loss:.4f}")

        # 5.1 Visualize Loss Curve
        loss_chart = VisualizerTools.generate_loss_chart(loss_history)
        VisualizerTools.render_chart_to_file(
            loss_chart,
            os.path.join(SRC_DIR, "exports/training_results/loss_curve.html"),
        )

        # 5.2 Save final model
        self.save_model()

        # 6. Final Evaluation on Validation Set
        if not val_set:
            logger.warning("⚠️ Validation set is empty. Skipping statistical analysis.")
            return

        logger.info(
            f"🧪 Final Evaluation: Base vs News-Integrated ({pred_len}-day Window)"
        )
        self.model.eval()
        predictor = KronosPredictor(self.model, self.tokenizer, device=self.device)

        base_maes = []
        news_maes = []

        print("\n" + "=" * 90)
        print(
            f"{'Date':<12} | {'Ticker':<8} | {'Base MAE':<15} | {'News MAE':<15} | {'Improvement'}"
        )
        print("-" * 90)

        for item in val_set:
            h = item["history"]
            t = item["target"]
            actuals = t["close"].values[:pred_len]

            x_ts = pd.to_datetime(h["date"])
            # Future timestamps: handle business days if possible, or just simple offset
            future_dates = pd.date_range(
                start=x_ts.iloc[-1] + timedelta(days=1), periods=pred_len, freq="B"
            )
            y_ts = pd.Series(future_dates)

            # A. Base Prediction
            p_base = predictor.predict(
                h, x_ts, y_ts, pred_len=pred_len, news_emb=None, verbose=False
            )
            b_preds = p_base["close"].values[: len(actuals)]

            # B. News-Aware Prediction
            p_news = predictor.predict(
                h,
                x_ts,
                y_ts,
                pred_len=pred_len,
                news_emb=item["news_emb"],
                verbose=False,
            )
            n_preds = p_news["close"].values[: len(actuals)]

            # Calculate MAE over the window
            b_mae = np.mean(np.abs(b_preds - actuals))
            n_mae = np.mean(np.abs(n_preds - actuals))

            base_maes.append(b_mae)
            news_maes.append(n_mae)

            improvement = (b_mae - n_mae) / (b_mae + 1e-6) * 100

            date_str = str(t["date"].values[0])[:10]
            ticker = h.iloc[-1]["ticker"] if "ticker" in h.columns else "Stock"
            print(
                f"{date_str:<12} | {ticker:<8} | {b_mae:<15.4f} | {n_mae:<15.4f} | {improvement:>+7.1f}%"
            )

            # C. Generate Visualization for this case
            try:
                # Helper to convert DF to KLinePoints
                def to_kp_list(preds_df):
                    points = []
                    for idx, row in preds_df.iterrows():
                        points.append(
                            KLinePoint(
                                date=str(idx)[:10],
                                open=row["open"],
                                high=row["high"],
                                low=row["low"],
                                close=row["close"],
                                volume=row["volume"] if "volume" in row else 0,
                            )
                        )
                    return points

                forecast_obj = ForecastResult(
                    ticker=ticker,
                    base_forecast=to_kp_list(p_base),
                    adjusted_forecast=to_kp_list(p_news),
                    rationale=item["summary"],
                )

                # Ground truth for visualizer expects a DataFrame with 'date' and 'close'
                gt_df = t[["date", "open", "high", "low", "close", "volume"]]

                chart = VisualizerTools.generate_stock_chart(
                    df=h,
                    ticker=ticker,
                    title=f"Training Eval: {ticker} ({date_str}) Improvement: {improvement:.1f}%",
                    forecast=forecast_obj,
                    ground_truth=gt_df,
                )

                safe_date = date_str.replace("-", "")
                filename = f"eval_{ticker}_{safe_date}.html"
                VisualizerTools.render_chart_to_file(
                    chart, os.path.join(SRC_DIR, f"exports/training_results/{filename}")
                )
            except Exception as e:
                logger.error(f"Failed to generate eval chart for {ticker}: {e}")

        # Summary Statistics
        avg_base_err = sum(base_maes) / max(1, len(base_maes))
        avg_news_err = sum(news_maes) / max(1, len(news_maes))
        overall_imp = (avg_base_err - avg_news_err) / (avg_base_err + 1e-6) * 100

        print("-" * 90)
        print(
            f"{'AVERAGE':<12} | {'-':<8} | {avg_base_err:<15.4f} | {avg_news_err:<15.4f} | {overall_imp:>+7.1f}%"
        )
        print("=" * 90 + "\n")

        logger.success(
            f"🏁 Statistical Analysis Complete. Avg Error Reduction ({pred_len}-day): {overall_imp:.2f}%"
        )
        logger.info(
            f"📊 Visualization results saved to: {os.path.join(SRC_DIR, 'exports/training_results/')}"
        )


if __name__ == "__main__":
    trainer = AutoSynthesisTrainer()

    logger.info("📂 Fetching all stock codes from database...")
    res = trainer.db.execute_query("SELECT code FROM stock_list")
    all_tickers = [row["code"] for row in res]

    if not all_tickers:
        logger.warning("⚠️ No tickers found in stock_list table. Trying to sync...")
        trainer.tools._check_and_update_stock_list(force=True)
        res = trainer.db.execute_query("SELECT code FROM stock_list")
        all_tickers = [row["code"] for row in res]

    logger.info(f"🚀 Starting training on potential stocks (1-year scan)...")
    # 为了演示，我们扫描前 100 个股票，寻找最近一年的冲击点
    trainer.run_synthesis_and_train(all_tickers[:100], pred_len=1)