# PRD 管理服务：上传、解析、版本快照

import json
import logging
import os
import time
import uuid
from datetime import datetime, timezone
from pathlib import Path

from server.config import settings
from server.services.prd_manager.word_parser import WordParser
from server.services.prd_manager.image_parser import ImageParser

logger = logging.getLogger("testflow")
RATE_LIMIT_DELAY = 0.5


class PRDService:
    """PRD management: handle file upload, text extraction, and version snapshots."""

    def __init__(self):
        self.output_dir = settings.OUTPUT_DIR
        os.makedirs(self.output_dir, exist_ok=True)
        self._prds: dict[str, dict] = {}

    async def upload_and_parse(self, filename: str, content: bytes) -> dict:
        prd_id = str(uuid.uuid4())[:8]
        basename = Path(filename).stem
        ext = Path(filename).suffix.lower()

        upload_dir = self.output_dir / prd_id
        os.makedirs(upload_dir, exist_ok=True)
        filepath = upload_dir / filename

        t0 = time.time()
        with open(filepath, "wb") as f:
            f.write(content)

        logger.info("[UPLOAD] 收到文件: %s (%.1f KB, %s)", filename, len(content) / 1024, ext)

        if ext in (".txt", ".md"):
            text = content.decode("utf-8", errors="replace")
            sections = self._parse_markdown_sections(text)
            parsed = {
                "source": str(filepath), "sections": sections,
                "image_sources": {}, "image_analysis": [], "full_text": text,
            }
            logger.info("[PARSE] %.4s %s → %d 章节, %d 字符, %.1fs",
                         prd_id, ext, len(sections), len(text), time.time() - t0)

        elif ext == ".docx":
            logger.info("[PARSE] %.4s → 开始解析 .docx...", prd_id)
            parsed = await self._parse_docx(str(filepath), upload_dir)
            logger.info("[PARSE] %.4s .docx 完成: %d 章节, %.1fs",
                         prd_id, len(parsed.get("sections", [])), time.time() - t0)

        elif ext == ".pdf":
            logger.info("[PARSE] %.4s → 开始解析 .pdf...", prd_id)
            parsed = await self._parse_pdf(str(filepath), upload_dir)
            logger.info("[PARSE] %.4s .pdf 完成: %d 页, %.1fs",
                         prd_id, len(parsed.get("sections", [])), time.time() - t0)
        else:
            raise ValueError(f"Unsupported file format: {ext}")

        parsed_path = upload_dir / f"{basename}_parsed.json"
        with open(parsed_path, "w", encoding="utf-8") as f:
            json.dump(parsed, f, ensure_ascii=False, indent=2)

        prd_version = {
            "prd_id": prd_id, "filename": filename,
            "uploaded_at": datetime.now(timezone.utc).isoformat(), "status": "ready",
            "parsed_path": str(parsed_path), "full_text": parsed.get("full_text", ""),
            "sections": parsed.get("sections", []), "images": parsed.get("image_analysis", []),
        }
        self._prds[prd_id] = prd_version

        n_paras = sum(1 for s in parsed.get("sections", []) for b in s.get("blocks", []) if b["type"] == "para")
        n_tables = sum(1 for s in parsed.get("sections", []) for b in s.get("blocks", []) if b["type"] == "table")
        logger.info("[UPLOAD] %.4s 解析完成: %d章节 %d段落 %d表格 %d图片 → %s",
                     prd_id, len(parsed.get("sections", [])), n_paras, n_tables,
                     len(parsed.get("image_analysis", [])), parsed_path)
        return prd_version

    async def upload_quick(self, filename: str, content: bytes) -> dict:
        """Quick save — returns immediately, parsing happens in background."""
        prd_id = str(uuid.uuid4())[:8]
        upload_dir = self.output_dir / prd_id
        os.makedirs(upload_dir, exist_ok=True)
        filepath = upload_dir / filename
        with open(filepath, "wb") as f:
            f.write(content)

        logger.info("[UPLOAD] %.4s 快速保存: %s (%.1f KB)", prd_id, filename, len(content) / 1024)

        prd_version = {
            "prd_id": prd_id, "filename": filename,
            "uploaded_at": datetime.now(timezone.utc).isoformat(),
            "status": "parsing", "full_text": "", "sections": [], "images": [],
            "parsed_path": "", "filepath": str(filepath),
        }
        self._prds[prd_id] = prd_version
        return prd_version

    async def parse_async(self, prd_id: str, filename: str, content: bytes):
        """Full parse (text + images + vision) — called in background thread."""
        basename = Path(filename).stem
        ext = Path(filename).suffix.lower()
        upload_dir = self.output_dir / prd_id
        filepath = upload_dir / filename
        t0 = time.time()

        # Parse
        if ext in (".txt", ".md"):
            text = content.decode("utf-8", errors="replace")
            sections = self._parse_markdown_sections(text)
            parsed = {"source": str(filepath), "sections": sections,
                      "image_sources": {}, "image_analysis": [], "full_text": text}
        elif ext == ".docx":
            parsed = await self._parse_docx(str(filepath), upload_dir)
        elif ext == ".pdf":
            parsed = await self._parse_pdf(str(filepath), upload_dir)
        else:
            self.set_status(prd_id, "error", f"Unsupported format: {ext}")
            return

        # Save parsed
        parsed_path = upload_dir / f"{basename}_parsed.json"
        with open(parsed_path, "w", encoding="utf-8") as f:
            json.dump(parsed, f, ensure_ascii=False, indent=2)

        # Update in-memory
        prd = self._prds.get(prd_id, {})
        prd.update({
            "status": "ready", "parsed_path": str(parsed_path),
            "full_text": parsed.get("full_text", ""),
            "sections": parsed.get("sections", []),
            "images": parsed.get("image_analysis", []),
        })

        n_paras = sum(1 for s in parsed.get("sections", []) for b in s.get("blocks", []) if b["type"] == "para")
        n_tables = sum(1 for s in parsed.get("sections", []) for b in s.get("blocks", []) if b["type"] == "table")
        logger.info("[UPLOAD] %.4s 后台解析完成: %d章节 %d段落 %d表格 %d图片, %.1fs",
                     prd_id, len(parsed.get("sections", [])), n_paras, n_tables,
                     len(parsed.get("image_analysis", [])), time.time() - t0)

    def set_status(self, prd_id: str, status: str, error: str = ""):
        if prd_id in self._prds:
            self._prds[prd_id]["status"] = status
            if error:
                self._prds[prd_id]["error"] = error

    async def get_status(self, prd_id: str) -> dict | None:
        prd = self._prds.get(prd_id)
        if not prd:
            return None
        return {"prd_id": prd_id, "status": prd.get("status", "ready"),
                "error": prd.get("error", ""),
                "sections_count": len(prd.get("sections", [])),
                "images_count": len(prd.get("images", []))}

    async def get_prd(self, prd_id: str) -> dict | None:
        return self._prds.get(prd_id)

    # ---- internal parsers ----

    async def _parse_docx(self, filepath: str, output_dir: Path) -> dict:
        word = WordParser(filepath)
        sections, image_sources = word.extract_sections()
        full_text = word.extract_full_text()

        images_dir = output_dir / "images"
        images = word.extract_images(str(images_dir))
        image_analysis = []
        n_chart_images = 0

        if images:
            parser = ImageParser()
            logger.info("[IMAGE] %.4s 发现 %d 张图片，开始视觉分析...", output_dir.parent.name, len(images))
            for i, img in enumerate(images):
                t_img = time.time()
                logger.info("[IMAGE] [%d/%d] rid=%s, 调用 Qwen VL...", i + 1, len(images), img["rid"])
                result = parser.parse_image(img["path"])
                if result is None:
                    result = {"type": "other", "description": ""}
                result["rid"] = img["rid"]
                result["path"] = img["path"]
                if img["rid"] in image_sources:
                    result["context"] = image_sources[img["rid"]]
                image_analysis.append(result)
                if result.get("type") in ("flowchart", "architecture", "state", "sequence", "activity"):
                    n_chart_images += 1
                logger.info("[IMAGE] [%d/%d] rid=%s 完成，类型=%s (%.1fs)",
                             i + 1, len(images), img["rid"], result.get("type", "?"), time.time() - t_img)
                if i < len(images) - 1:
                    time.sleep(RATE_LIMIT_DELAY)

            usg = parser.usage
            logger.info("[IMAGE] %.4s 图片分析完成: %d张, 图表类%d张, tokens(p=%d c=%d t=%d)",
                         output_dir.parent.name, len(images), n_chart_images,
                         usg["prompt_tokens"], usg["completion_tokens"], usg["total_tokens"])
        else:
            logger.info("[IMAGE] %.4s 文档无图片", output_dir.parent.name)

        return {
            "source": filepath, "sections": sections,
            "image_sources": image_sources, "image_analysis": image_analysis,
            "full_text": full_text,
        }

    def _parse_markdown_sections(self, text: str) -> list[dict]:
        import re
        sections: list[dict] = []
        current_source = "正文"
        current_blocks: list[dict] = []
        para_idx = 0

        for line in text.split("\n"):
            stripped = line.strip()
            if not stripped:
                continue
            heading_match = re.match(r"^(#{1,6})\s+(.+)", stripped)
            if heading_match:
                if current_blocks:
                    sections.append({"source": current_source, "blocks": current_blocks, "images": []})
                    current_blocks = []
                level = len(heading_match.group(1))
                current_source = f"{'#' * level} {heading_match.group(2)}"
                continue
            if "|" in stripped and stripped.count("|") >= 2:
                cells = [c.strip() for c in stripped.split("|") if c.strip()]
                if cells:
                    current_blocks.append({
                        "type": "table",
                        "table": len([b for b in current_blocks if b["type"] == "table"]) + 1,
                        "headers": cells,
                        "rows": [{"columns": [{"name": c, "row": 1, "col": i + 1, "text": c} for i, c in enumerate(cells)]}],
                    })
                continue
            para_idx += 1
            current_blocks.append({"type": "para", "index": para_idx, "text": stripped})

        if current_blocks:
            sections.append({"source": current_source, "blocks": current_blocks, "images": []})
        return sections

    async def _parse_pdf(self, filepath: str, output_dir: Path) -> dict:
        from PyPDF2 import PdfReader
        reader = PdfReader(filepath)
        full_text_parts = []
        for i, page in enumerate(reader.pages):
            text = page.extract_text()
            if text:
                full_text_parts.append(text)
        full_text = "\n".join(full_text_parts)
        sections = [{"source": "PDF 内容", "blocks": [{"type": "para", "index": i + 1, "text": t} for i, t in enumerate(full_text_parts)], "images": []}]
        return {"source": filepath, "sections": sections, "image_sources": {}, "image_analysis": [], "full_text": full_text}


prd_service = PRDService()