# -*- coding: utf-8 -*-
"""中文写作辅助系统 v3.3（教学增强版）"""
import os
import re
import json
import uuid
import difflib
import asyncio
import sqlite3
import argparse
import random
from datetime import datetime, timedelta
from contextlib import contextmanager
from filelock import FileLock
from dotenv import load_dotenv
from rich.console import Console
from rich.markup import escape
from aiohttp import web, ClientSession

# ================ 初始化配置 ================
load_dotenv()
console = Console(record=True)

# ================ 常量定义 ================
API_CONFIG = {
    "DEEPSEEK": {
        "api_key": os.getenv("DEEPSEEK_API_KEY"),
        "model": "deepseek-chat",
        "endpoint": "https://api.deepseek.com/v1/chat/completions"
    }
}

# 检查 API 配置
if not API_CONFIG["DEEPSEEK"]["api_key"]:
    console.print("[red]错误：未设置 DEEPSEEK_API_KEY 环境变量[/]")
    exit(1)

# ================ 数据库模块 ================
class ConversationDB:
    # 保持原有实现不变
    def __init__(self, db_path='conversations.db'):
        self.db_path = db_path
        self._init_db()

    @contextmanager
    def _get_connection(self):
        conn = sqlite3.connect(self.db_path)
        conn.row_factory = sqlite3.Row
        try:
            yield conn
        finally:
            conn.close()

    def _init_db(self):
        with self._get_connection() as conn:
            conn.execute('''
                CREATE TABLE IF NOT EXISTS conversations (
                    session_id TEXT PRIMARY KEY,
                    history TEXT NOT NULL,
                    last_accessed TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
            ''')
            conn.commit()

    async def save_conversation(self, session_id, history):
        loop = asyncio.get_running_loop()
        await loop.run_in_executor(
            None, 
            self._sync_save, 
            session_id, 
            json.dumps(history, ensure_ascii=False)
        )

    def _sync_save(self, session_id, history_json):
        with self._get_connection() as conn:
            conn.execute('''
                INSERT OR REPLACE INTO conversations 
                VALUES (?, ?, ?)
            ''', (session_id, history_json, datetime.now()))
            conn.commit()

    async def load_conversation(self, session_id):
        loop = asyncio.get_running_loop()
        return await loop.run_in_executor(
            None,
            self._sync_load,
            session_id
        )

    def _sync_load(self, session_id):
        with self._get_connection() as conn:
            cursor = conn.execute('''
                SELECT history FROM conversations 
                WHERE session_id = ?
            ''', (session_id,))
            result = cursor.fetchone()
            return json.loads(result['history']) if result else None

    async def cleanup_old_sessions(self, max_age_days=7):
        loop = asyncio.get_running_loop()
        await loop.run_in_executor(
            None,
            self._sync_cleanup,
            max_age_days
        )

    def _sync_cleanup(self, max_age_days):
        cutoff = datetime.now() - timedelta(days=max_age_days)
        with self._get_connection() as conn:
            conn.execute('''
                DELETE FROM conversations 
                WHERE last_accessed < ?
            ''', (cutoff,))
            conn.commit()

# ================ 工具函数 ================
def highlight_text(text: str) -> str:
    text = escape(text)
    patterns = [
        (r"\{\~\~(.*?)\~\~\}", "[strike red]\\1[/]"),
        (r"\{\+\+(.*?)\+\+\}", "[underline green]\\1[/]")
    ]
    for pattern, replacement in patterns:
        text = re.sub(pattern, replacement, text)
    return text

def load_training_examples(file_path: str) -> str:
    with open(file_path, "r", encoding="utf-8") as f:
        data = json.load(f)
    
    prompt_lines = ["**学习示例**："]
    for category in data:
        if isinstance(data[category], dict):
            for sub_type in data[category]:
                prompt_lines.append(f"\n▌ {category}-{sub_type}：")
                for idx, example in enumerate(data[category][sub_type][:3], 1):
                    prompt_lines.append(f"{idx}. 错误：{example['原句']}")
                    prompt_lines.append(f"   修正：{example['修正']}")
        else:
            prompt_lines.append(f"\n▌ {category}：")
            for idx, example in enumerate(data[category][:3], 1):
                prompt_lines.append(f"{idx}. 错误：{example['原句']}")
                prompt_lines.append(f"   修正：{example['修正']}")
    
    return "\n".join(prompt_lines)

def extract_error_types(response: str) -> list:
    """增强错误类型提取（支持多类型分割）"""
    matches = re.findall(r"错误类型：([^\n]+)", response)
    return list(set([et.strip() for m in matches for et in m.split("、") if et.strip()]))

def save_to_training(data: dict, error_type: str):
    main_type, sub_type = error_type.split("-", 1)
    with FileLock("training_data.json.lock"):
        with open("training_data.json", "r+", encoding="utf-8") as f:
            training = json.load(f)
            if main_type not in training:
                training[main_type] = {}
            if sub_type not in training[main_type]:
                training[main_type][sub_type] = []
            
            if not any(ex["原句"] == data["原句"] for ex in training[main_type][sub_type]):
                training[main_type][sub_type].append(data)
                f.seek(0)
                json.dump(training, f, ensure_ascii=False, indent=2)
                f.truncate()

# ================ 核心处理模块 ================
class WritingAssistant:
    def __init__(self, test_mode=False):
        self.test_mode = test_mode
        self.client_session = None
        self.lock = asyncio.Lock()
        self.system_prompt = self._build_system_prompt()
        self.exercises = {}
        self.conversations = {} if test_mode else ConversationDB()
        asyncio.create_task(self.auto_cleanup_task())

    def _build_system_prompt(self):
        training_content = load_training_examples("training_data.json")
        return f"""
作为专业中文写作修改专家，请严格遵循以下规范：

{training_content}

**修改规则**：
1. 必须使用 {{~~删除内容~~}} 和 {{++新增内容++}} 标记
2. 仅修正语法错误，保持原意不变
3. 必须标注错误类型（格式：错误类型：类型名称）
"""

    async def auto_cleanup_task(self, interval_hours=24, max_age_days=7):
        while True:
            await asyncio.sleep(interval_hours * 3600)
            if not self.test_mode:
                await self.conversations.cleanup_old_sessions(max_age_days)
            expired = [k for k in self.exercises if 
                      (datetime.now() - self.exercises[k]["timestamp"]).days > 7]
            for key in expired:
                del self.exercises[key]
            console.log(f"已完成自动清理")

    async def ensure_session(self):
        if self.client_session is None or self.client_session.closed:
            self.client_session = ClientSession(headers={"Content-Type": "application/json"})

    async def analyze(self, text: str, session_id: str) -> str:
        async with self.lock:
            history = await self._get_history(session_id)
            history.append({"role": "user", "content": text})

            try:
                await self.ensure_session()
                raw_response = await self._call_deepseek(history)
                await self._process_exercises(raw_response, session_id)
                await self._save_history(session_id, history + [{"role": "assistant", "content": raw_response}])
                return raw_response
            except Exception as e:
                if len(history) > 1:
                    history.pop()
                console.print(f"[red]错误：{str(e)}[/]")
                raise

    async def _get_history(self, session_id):
        if self.test_mode:
            return self.conversations.get(session_id, [{"role": "system", "content": self.system_prompt}])
        else:
            return await self.conversations.load_conversation(session_id) or [{"role": "system", "content": self.system_prompt}]

    async def _save_history(self, session_id, history):
        if self.test_mode:
            self.conversations[session_id] = history
        else:
            await self.conversations.save_conversation(session_id, history)

    def _get_preset_examples(self, error_type: str, count: int) -> list:
        """修复预设样例获取逻辑"""
        try:
            if "-" not in error_type:
                return []
            main_type, sub_type = error_type.split("-", 1)
            with open("training_data.json", "r", encoding="utf-8") as f:
                data = json.load(f)
                candidates = data.get(main_type, {}).get(sub_type, [])
                valid_samples = [ex for ex in candidates if "原句" in ex and ex["原句"] != "生成失败示例"]
                return random.sample(valid_samples, min(count, len(valid_samples)))
        except Exception as e:
            console.log(f"预设样例获取失败: {str(e)}")
            return []

    async def _generate_new_example(self, error_type: str) -> dict:
        """增强JSON生成与解析"""
        prompt = f"""请生成包含【{error_type}】错误的句子，按以下严格的JSON格式响应：
        {{
            "原句": "包含错误的句子（必须包含中文双引号）",
            "修正": "修正后的句子（必须包含中文双引号）",
            "说明": "错误类型说明"
        }}
        不要包含任何额外文本和代码块标记"""
        
        try:
            response = await self._call_deepseek(
                [{"role": "user", "content": prompt}],
                temperature=0.7
            )
            # 增强JSON提取
            json_str = re.search(r'\{[\s\S]*\}', response).group()
            return json.loads(json_str)
        except Exception as e:
            console.log(f"生成失败：{str(e)}")
            return {"原句": "生成失败示例", "修正": "请重试生成"}

    async def _process_exercises(self, response, session_id):
        """改进的多错误类型处理流程"""
        try:
            error_types = extract_error_types(response)
            if not error_types:
                return

            exercises = {}
            for et in error_types:
                if "-" not in et:
                    continue

                # 独立处理每个错误类型
                preset = self._get_preset_examples(et, 2)
                generated = await self._generate_new_example(et)

                # 独立保存每个类型的生成样例
                if "生成失败示例" not in generated["原句"]:
                    await self._save_training_data(generated, et)  # 同步保存

                # 为每个类型创建独立练习组
                exercises[et] = {
                    "preset": preset[:2],
                    "generated": generated,
                    "timestamp": datetime.now()
                }

            self.exercises[session_id] = exercises

        except Exception as e:
            console.log(f"练习生成失败: {str(e)}")

    async def _save_training_data(self, data, error_type):
        """增强的数据保存逻辑"""
        try:
            main_type, sub_type = error_type.split("-", 1)
            with FileLock("training_data.json.lock"):
                with open("training_data.json", "r+", encoding="utf-8") as f:
                    training = json.load(f)

                    # 确保分类结构存在
                    training.setdefault(main_type, {}).setdefault(sub_type, [])

                    # 检查重复性
                    existing = {ex["原句"] for ex in training[main_type][sub_type]}
                    if data["原句"] not in existing:
                        training[main_type][sub_type].append(data)
                        f.seek(0)
                        json.dump(training, f, ensure_ascii=False, indent=2)
                        f.truncate()

        except Exception as e:
            console.log(f"保存失败: {str(e)}")

    async def _call_deepseek(self, messages, temperature=0.1):
        payload = {
            "messages": messages,
            "model": API_CONFIG["DEEPSEEK"]["model"],
            "temperature": temperature
        }
        
        try:
            async with self.client_session.post(
                API_CONFIG["DEEPSEEK"]["endpoint"],
                headers={
                    "Authorization": f"Bearer {API_CONFIG['DEEPSEEK']['api_key']}",
                    "Content-Type": "application/json"
                },
                json=payload,
                timeout=30
            ) as resp:
                if resp.status != 200:
                    error_text = await resp.text()
                    raise ValueError(f"API 请求失败 (状态码: {resp.status}): {error_text}")
                
                data = await resp.json()
                if not data.get("choices") or not data["choices"][0].get("message", {}).get("content"):
                    raise ValueError("API 响应格式不正确")
                return data["choices"][0]["message"]["content"]
        except Exception as e:
            console.print(f"[red]错误：API 调用失败 - {str(e)}[/]")
            raise

# ================ HTTP服务模块 ================
async def handle_post(request):
    try:
        data = await request.json()
        input_text = data.get("input")
        session_id = data.get("session_id", str(uuid.uuid4()))

        if not input_text:
            return web.json_response({"code": 1, "out": "缺少输入内容"}, status=400)

        assistant = request.app["assistant"]
        analysis_result = await assistant.analyze(input_text, session_id)
        
        # 获取分类型练习数据
        exercises_data = assistant.exercises.get(session_id, {})

        # 构建分类型响应结构
        structured_response = {
            "analysis": analysis_result,
            "exercises": [
                {
                    "error_type": et,
                    "preset_examples": [
                        {"original": ex["原句"], "corrected": ex["修正"]}
                        for ex in ex_data["preset"]
                    ],
                    "generated_example": {
                        "original": ex_data["generated"]["原句"],
                        "corrected": ex_data["generated"]["修正"]
                    }
                }
                for et, ex_data in exercises_data.items()
            ]
        }

        return web.json_response({
            "code": 0,
            "data": structured_response,
            "session_id": session_id
        })
    
    except Exception as e:
        return web.json_response({"code": 2, "out": str(e)})

async def handle_exercise(request):
    try:
        data = await request.json()
        session_id = data["session_id"]
        answers = data.get("answers", "")
        
        user_answers = {}
        for line in answers.split("\n"):
            if "=>" in line:
                parts = line.split("=>", 1)
                user_answers[parts[0].strip()] = parts[1].strip()
        
        assistant = request.app["assistant"]
        exercises = assistant.exercises.get(session_id, {}).get("data", {})
        
        results = []
        for error_type, cases in exercises.items():
            for case in cases:
                std = case["修正"]
                user = user_answers.get(case["原句"], "未提交")
                
                seq = difflib.SequenceMatcher(None, std, user)
                ratio = seq.ratio()
                
                results.append({
                    "error_type": error_type,
                    "original": case["原句"],
                    "user_answer": user,
                    "std_answer": std,
                    "similarity": f"{ratio*100:.1f}%"
                })
        
        return web.json_response({
            "code": 0,
            "results": results
        })
    
    except Exception as e:
        return web.json_response({"code": 3, "error": str(e)})

@web.middleware
async def cors_middleware(request, handler):
    response = await handler(request)
    response.headers["Access-Control-Allow-Origin"] = "*"
    response.headers["Access-Control-Allow-Methods"] = "POST, OPTIONS"
    response.headers["Access-Control-Allow-Headers"] = "Content-Type"
    return response

async def init_app():
    app = web.Application(middlewares=[cors_middleware])
    app["assistant"] = WritingAssistant()
    app.router.add_post("/api/", handle_post)
    app.router.add_post("/api/exercise", handle_exercise)
    app.router.add_options("/api/", lambda _: web.Response())
    return app

# ================ 测试模块 ================
async def run_interactive_test():
    assistant = WritingAssistant(test_mode=True)
    session_id = "interactive_session"
    console.print("\n[bold magenta]=== 进入交互测试模式 ===")
    console.print("[italic]输入'exit'退出，输入内容后按回车发送[/]\n")

    while True:
        try:
            user_input = console.input("[bold cyan]用户输入 > [/]")
            if user_input.lower() in ['exit', 'quit']:
                console.print("[yellow]退出交互测试模式[/]")
                break
                
            response = await assistant.analyze(user_input, session_id)
            
            console.print(f"\n[bold green]模型分析：[/]{highlight_text(response)}")
            
            
            
            console.print("[bold white]" + "-"*50 + "[/]\n")
            
        except KeyboardInterrupt:
            console.print("\n[yellow]检测到中断信号，退出测试[/]")
            break
        except Exception as e:
            break

def run_production_mode():
    web.run_app(init_app(), port=2346)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="中文写作辅助系统")
    parser.add_argument("--test", action="store_true", help="交互测试模式")
    args = parser.parse_args()

    if args.test:
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
        loop.run_until_complete(run_interactive_test())
    else:
        run_production_mode()