import os
from unsloth import FastLanguageModel
import torch
from trl import SFTTrainer
from transformers import TrainingArguments
from datasets import load_dataset

# 加载模型
max_seq_length = 2048
dtype = None
load_in_4bit = True
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name = "/home/zk/ai/base_model/llama-3-8b-bnb-4bit",
    max_seq_length = max_seq_length,
    dtype = dtype,
    load_in_4bit = load_in_4bit,
)

# 准备训练数据
alpaca_prompt = """Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.
### Instruction：
{}
### Input:
{}
### Response:
{}"""

EOS_TOKEN = tokenizer.eos_token  # 必须添加 EOS_TOKEN
def formatting_prompts_func(examples):
    instructions = examples["instruction"]
    inputs = examples["input"]
    outputs = examples["output"]
    texts = []
    for instruction, input, output in zip(instructions, inputs, outputs):
        # 必须添加EOS_TOKEN，否则无限生成
        text = alpaca_prompt.format(instruction, input, output) + EOS_TOKEN
        texts.append(text)
    return {"text": texts}
pass

# 本地数据集
dataset = load_dataset("json", data_files = "/home/zk/ai/dataset/caishui_2011_100hao.json", split="train")
dataset = dataset.map(formatting_prompts_func, batched = True)

# 设置训练参数
model = FastLanguageModel.get_peft_model(
    model,
    r = 16,
    target_modules = ["q_proj", "k_proj", "v_proj", "o_proj",
                        "gate_proj", "up_proj", "down_proj", ],
    lora_alpha = 16,
    lora_dropout = 0,
    bias = "none",
    use_gradient_checkpointing = True,
    random_state = 3407,
    max_seq_length = max_seq_length,
    use_rslora = False,
    loftq_config = None,
)

trainer = SFTTrainer(
    model = model,
    train_dataset = dataset,
    dataset_text_field = "text",
    max_seq_length = max_seq_length,
    tokenizer = tokenizer,
    args = TrainingArguments(
        per_device_train_batch_size = 1,
        gradient_accumulation_steps = 4,
        warmup_steps = 2,
        max_steps = 20,
        fp16 = not torch.cuda.is_bf16_supported(),
        bf16 = torch.cuda.is_bf16_supported(),
        logging_steps = 1,
        output_dir = "outputs",
        optim = "adamw_8bit",
        weight_decay = 0.01,
        lr_scheduler_type = "linear",
        seed = 3407,
        learning_rate = 2e-5,
    ),
)

# 开始训练
trainer.train()

# 保存微调模型
model.save_pretrained("lora_model")

# 选项：保存为16位hf模型
save_16bit = input("是否保存为16位hf模型？(y/n): ")
if save_16bit.lower() == "y":
    model.save_pretrained_merged("outputs", tokenizer, save_method="merged_16bit")

# 选项：保存为gguf模型
save_gguf = input("是否保存为gguf模型？(y/n): ")
if save_gguf.lower() == "y":
    os.system("python /home/zk/ai/llama.cpp/convert_hf_to_gguf.py --outfile /home/zk/ai/gguf_model/lm38b_tax_jzjt.gguf /home/zk/ai/outputs")

# 选项：量化为4位gguf模型
quantize_4bit = input("是否量化为4位gguf模型？(y/n): ")
if quantize_4bit.lower() == "y":
    os.system("/home/zk/ai/llama.cpp/build/bin/llama-quantize /home/zk/ai/gguf_model/QWQ_tax_jzjt.gguf /home/zk/ai/gguf_model/lm38b_tax_jzjt-Q4_K_M.gguf Q4_K_M")

def split_text_into_parts(text_content, part_length=3000):
    """智能文本分段，避免截断完整句子"""
    parts = []
    sentences = re.split(r'(?<=[。！？!?])', text_content)
    
    current_part = ""
    for sentence in sentences:
        if len(current_part) + len(sentence) <= part_length:
            current_part += sentence
        else:
            if current_part:
                parts.append(current_part)
            current_part = sentence
    
    if current_part:
        parts.append(current_part)
    
    return parts

def extract_movie_name(file_path):
    """从复杂文件名中智能提取电影名称"""
    filename = Path(file_path).stem
    
    # 移除下载站点标记
    clean_name = re.sub(r'^\[.*?\]', '', filename)
    
    # 移除分辨率、编码等信息
    patterns_to_remove = [
        r'\d{4}p',  # 1080p, 720p等
        r'BluRay', 'DVD', 'WEB-DL',
        r'x\d{3}', 'H\.?264', 'HEVC',
        r'\d{4}\.\d{2}\.\d{2}',  # 日期
        r'[A-Z][A-Za-z]+-?Team',  # 发布组
    ]
    
    for pattern in patterns_to_remove:
        clean_name = re.sub(pattern, '', clean_name, flags=re.IGNORECASE)
    
    # 清理多余字符
    clean_name = re.sub(r'[_-]+', ' ', clean_name)
    clean_name = re.sub(r'\s+', ' ', clean_name).strip()
    
    return clean_name

class ConversationManager:
    """管理多轮对话上下文"""
    def __init__(self, max_context_tokens=4000):
        self.history = []
        self.max_tokens = max_context_tokens
        
    def add_message(self, role, content):
        """添加消息并维护上下文长度"""
        self.history.append({"role": role, "content": content})
        self._trim_history()
    
    def _trim_history(self):
        """智能修剪历史记录，保留最重要的对话"""
        total_tokens = sum(len(msg["content"]) for msg in self.history)
        
        while total_tokens > self.max_tokens and len(self.history) > 2:
            # 移除最旧的非系统消息
            for i, msg in enumerate(self.history):
                if msg["role"] != "system":
                    removed = self.history.pop(i)
                    total_tokens -= len(removed["content"])
                    break

# 第一阶段：告知模型准备接收长内容
initial_prompt = f"""我将分{parts_count}次发送电影《{movie_name}》的完整台词给你。
请你先接收所有台词内容，不要立即分析。
等我发送完所有部分后，我会让你开始分析。

这是第1部分，共{parts_count}部分：
{first_part}

请回复"收到第1部分，等待后续内容。"
"""

analysis_prompt = f"""请根据所有台词内容完成以下分析，严格按照要求的格式输出：

电影名称：{movie_name}
经典台词：
1. [第一句经典台词]
2. [第二句经典台词]
...（必须输出10句）

微小说：[根据台词创作的300字微小说。]

要求（重要指令）：
1. 根据你所知道的给出该电影超级简单的介绍
2. 经典台词要选择最具代表性、最打动人心的10句
3. 微小说要基于所有台词的意境和主题创作，保持故事完整性
4. 注意：不要和原电影叙事年代相同，尽量用现在的年代
5. 记得为微小说写题目
6. 禁止：不要输出任何评价性内容，只输出要求的格式内容

严格按照上述格式，不要添加额外说明！"""

system_prompt = """你是一位资深的电影评论家和小说作家，具有以下特点：
1. 对电影艺术有深刻理解，能准确把握电影主题
2. 擅长从台词中提炼精华，识别经典对白
3. 具有出色的改编能力，能将经典故事现代化
4. 严格遵守格式要求，输出结构清晰

你的任务是根据提供的电影台词，完成专业分析。"""

def create_staged_prompts(movie_name, parts):
    """创建分阶段的prompt序列"""
    prompts = []
    
    # 阶段1：建立上下文
    prompts.append({
        "role": "system",
        "content": f"你将接收电影《{movie_name}》的完整台词，请先专心接收，不要分析。"
    })
    
    # 阶段2：分段发送内容
    for i, part in enumerate(parts, 1):
        if i == 1:
            content = f"开始发送《{movie_name}》台词，这是第{i}部分：\n\n{part}\n\n请确认收到。"
        else:
            content = f"继续发送《{movie_name}》台词，第{i}部分：\n\n{part}\n\n请确认收到。"
        prompts.append({"role": "user", "content": content})
    
    # 阶段3：分析指令
    prompts.append({
        "role": "user", 
        "content": """现在你已经接收了完整台词，请开始分析。注意以下关键要求：
        1. 输出10句经典台词，必须编号1-10
        2. 微小说要现代背景，300字左右
        3. 不要评价，只输出结果"""
    })
    
    return prompts

# 简单三步配置
API_KEY = "您的API密钥"
INPUT_FILE_PATH = "电影台词.txt"
MODEL_NAME = "deepseek-ai/DeepSeek-R1"

def clean_srt_content(content):
    """清理SRT字幕文件内容"""
    lines = content.split('\n')
    cleaned_lines = []
    i = 0
    n = len(lines)
    
    while i < n:
        line = lines[i].strip()
        
        # 跳过序号行和时间码行
        if line.isdigit() or '-->' in line:
            i += 1
            continue
            
        # 保留并清理台词行
        if line:
            line = re.sub(r'\{[^}]*\}', '', line)
            line = line.strip()
            if line:
                cleaned_lines.append(line)
        
        i += 1
    
    return '\n'.join(cleaned_lines)

# 创建下载器实例
downloader = SubtitleDownloader(cookie)

# 下载日期页面字幕
downloader.download_from_date_pages(start_year=2024, start_month=3, end_month=12)

# 下载顺序页面字幕
downloader.download_from_sequence_pages(start_id=151336, end_id=151400)

# 下载2024年全年的字幕
downloader.download_from_date_pages(start_year=2024, start_month=1, end_month=12)

# 下载大量顺序页面
downloader.download_from_sequence_pages(start_id=150000, end_id=160000)

# 在parse_page()方法中添加调试信息
print(f"页面响应：{response.text[:500]}")