import json
import asyncio
import logging
import time
import uuid
from dataclasses import dataclass, field
from typing import List, Optional, Dict

import httpx

from ..core.config import settings
from ..core.model_resolver import get_active_embedding_config
from .chunk_engine import ChunkEngine, DocumentChunk, ChunkResult, _detect_csv_type, csv_to_markdown

logger = logging.getLogger(__name__)

EMBEDDING_BATCH_SIZE = 20
EMBEDDING_INTERVAL_SEC = 1
EMBEDDING_MAX_RETRIES = 3
EMBEDDING_RETRY_BACKOFF = 2

EMBEDDING_MAX_INPUT_CHARS = 12000
EMBEDDING_PAYLOAD_MAX_CHARS = 12000

GRAPHITI_CONCURRENCY = 5
GRAPHITI_TIMEOUT_SEC = 60

QDRANT_UUID_NAMESPACE = uuid.UUID("6ba7b810-9dad-11d1-80b4-00c04fd430c8")


def _chunk_id_to_uuid(chunk_id: str) -> str:
    return str(uuid.uuid5(QDRANT_UUID_NAMESPACE, chunk_id))


@dataclass
class ImportOptions:
    target_collection: str = ""
    import_qdrant: bool = False
    import_graphiti: bool = False
    chunk_strategy: str = "auto"
    concurrency: int = 3
    csv_content: Optional[str] = None
    region: str = ""
    issuer: str = ""
    effective_date: str = ""
    doc_version: str = ""


@dataclass
class ImportResult:
    success: bool = False
    chunk_info: Dict = field(default_factory=dict)
    qdrant_result: Optional[Dict] = None
    graphiti_result: Optional[Dict] = None
    errors: List[str] = field(default_factory=list)


class QdrantChunkImporter:
    def __init__(self):
        self._embed_config: Optional[Dict] = None
        self.qdrant_host = settings.QDRANT_HOST
        self.qdrant_port = settings.QDRANT_PORT
        self.doc_metadata: Dict = {}

    async def _ensure_config(self):
        if self._embed_config is None:
            self._embed_config = await get_active_embedding_config()
            logger.info(f"Embedding 模型配置已加载: {self._embed_config.get('model_name')} @ {self._embed_config.get('api_base')}")
        return self._embed_config

    @property
    def embed_model(self):
        return (self._embed_config or {}).get("model_name", "")

    @property
    def embed_api_base(self):
        return (self._embed_config or {}).get("api_base", "")

    @property
    def embed_api_key(self):
        return (self._embed_config or {}).get("api_key", "")

    async def import_chunk(self, chunk: DocumentChunk, collection: str) -> dict:
        embedding_input = self._build_embedding_input(chunk)
        vector = await self._get_embedding(embedding_input)
        if vector is None:
            return {"status": "failed", "error": "Embedding failed", "chunk_id": chunk.chunk_id}

        payload = self._build_payload(chunk)
        point_id = _chunk_id_to_uuid(chunk.chunk_id)
        payload["title"] = chunk.chunk_id

        try:
            async with httpx.AsyncClient(timeout=30) as client:
                res = await client.put(
                    f"http://{self.qdrant_host}:{self.qdrant_port}/collections/{collection}/points",
                    json={"points": [{"id": point_id, "vector": vector, "payload": payload}]},
                )
                if res.status_code in (200, 201):
                    return {"status": "done", "point_id": point_id}
                else:
                    return {"status": "failed", "error": res.text, "chunk_id": chunk.chunk_id}
        except Exception as e:
            return {"status": "failed", "error": str(e), "chunk_id": chunk.chunk_id}

    async def import_chunks_batch(self, chunks: List[DocumentChunk], collection: str) -> dict:
        total = len(chunks)
        success_count = 0
        failed_count = 0
        point_ids = []
        errors = []

        for batch_start in range(0, total, EMBEDDING_BATCH_SIZE):
            batch = chunks[batch_start:batch_start + EMBEDDING_BATCH_SIZE]
            texts = [self._build_embedding_input(c) for c in batch]

            vectors = await self._get_embeddings_batch(texts)
            if vectors is None:
                for i, (chunk, text) in enumerate(zip(batch, texts)):
                    vector = await self._get_embedding_single(text)
                    if vector:
                        pid = _chunk_id_to_uuid(chunk.chunk_id)
                        payload = self._build_payload(chunk)
                        payload["title"] = chunk.chunk_id
                        ok = await self._upsert_point(collection, pid, vector, payload)
                        if ok:
                            success_count += 1
                            point_ids.append(pid)
                        else:
                            failed_count += 1
                            errors.append(f"Qdrant upsert failed: {chunk.chunk_id}")
                    else:
                        failed_count += 1
                        errors.append(f"Embedding failed for chunk: {chunk.chunk_id}")
                if batch_start + EMBEDDING_BATCH_SIZE < total:
                    await asyncio.sleep(EMBEDDING_INTERVAL_SEC)
                continue

            points = []
            for chunk, vector in zip(batch, vectors):
                payload = self._build_payload(chunk)
                point_id = _chunk_id_to_uuid(chunk.chunk_id)
                payload["title"] = chunk.chunk_id
                points.append({
                    "id": point_id,
                    "vector": vector,
                    "payload": payload,
                })

            try:
                async with httpx.AsyncClient(timeout=60) as client:
                    res = await client.put(
                        f"http://{self.qdrant_host}:{self.qdrant_port}/collections/{collection}/points",
                        json={"points": points},
                    )
                    if res.status_code in (200, 201):
                        success_count += len(points)
                        point_ids.extend([p["id"] for p in points])
                    else:
                        failed_count += len(points)
                        errors.append(f"Qdrant upsert failed: {res.text[:200]}")
            except Exception as e:
                failed_count += len(points)
                errors.append(f"Qdrant upsert error: {str(e)}")

            if batch_start + EMBEDDING_BATCH_SIZE < total:
                await asyncio.sleep(EMBEDDING_INTERVAL_SEC)

        return {
            "status": "done" if failed_count == 0 else "partial",
            "total": total,
            "success_count": success_count,
            "failed_count": failed_count,
            "point_ids": point_ids,
            "errors": errors[:10],
        }

    async def _get_embedding_single(self, text: str) -> Optional[List[float]]:
        await self._ensure_config()
        for attempt in range(EMBEDDING_MAX_RETRIES):
            try:
                async with httpx.AsyncClient(timeout=60) as client:
                    res = await client.post(
                        f"{self.embed_api_base}/embeddings",
                        headers={
                            "Authorization": f"Bearer {self.embed_api_key}",
                            "Content-Type": "application/json",
                        },
                        json={"model": self.embed_model, "input": [text[:EMBEDDING_MAX_INPUT_CHARS]], "encoding_format": "float"},
                    )
                    if res.status_code == 200:
                        data = res.json()["data"]
                        return data[0]["embedding"]
                    elif res.status_code == 429:
                        await asyncio.sleep(EMBEDDING_RETRY_BACKOFF ** attempt)
                    else:
                        logger.error(f"Embedding single API error: {res.status_code} {res.text[:200]}")
                        return None
            except Exception as e:
                logger.error(f"Embedding single API exception: {e}")
                if attempt < EMBEDDING_MAX_RETRIES - 1:
                    await asyncio.sleep(EMBEDDING_RETRY_BACKOFF ** attempt)
        return None

    async def _upsert_point(self, collection: str, point_id: str, vector: List[float], payload: dict) -> bool:
        try:
            async with httpx.AsyncClient(timeout=30) as client:
                res = await client.put(
                    f"http://{self.qdrant_host}:{self.qdrant_port}/collections/{collection}/points",
                    json={"points": [{"id": point_id, "vector": vector, "payload": payload}]},
                )
                return res.status_code in (200, 201)
        except Exception:
            return False

    def _build_embedding_input(self, chunk: DocumentChunk) -> str:
        if chunk.source_type == "laws":
            return chunk.content[:EMBEDDING_MAX_INPUT_CHARS] if chunk.content else chunk.full_text[:EMBEDDING_MAX_INPUT_CHARS]
        elif chunk.source_type == "standards":
            parts = [chunk.violation, chunk.legal_basis]
            if chunk.severity_levels:
                parts.append("; ".join(f"{s.level}:{s.penalty}" for s in chunk.severity_levels))
            if chunk.applicable_conditions:
                parts.append(chunk.applicable_conditions[:500])
            text = " ".join(p for p in parts if p)
            return text if text else chunk.full_text[:EMBEDDING_MAX_INPUT_CHARS]
        elif chunk.source_type == "cases":
            parts = [chunk.violation_type]
            if chunk.legal_basis_list:
                parts.extend(chunk.legal_basis_list)
            parts.append(chunk.penalty_summary)
            text = " ".join(p for p in parts if p)
            return text if text else chunk.full_text[:EMBEDDING_MAX_INPUT_CHARS]
        else:
            return chunk.content[:EMBEDDING_MAX_INPUT_CHARS] if chunk.content else chunk.full_text[:EMBEDDING_MAX_INPUT_CHARS]

    def _build_payload(self, chunk: DocumentChunk) -> dict:
        base_meta = {}
        if self.doc_metadata:
            if self.doc_metadata.get("region"):
                base_meta["region"] = self.doc_metadata["region"]
            if self.doc_metadata.get("issuer"):
                base_meta["issuer"] = self.doc_metadata["issuer"]
            if self.doc_metadata.get("effective_date"):
                base_meta["effective_date"] = self.doc_metadata["effective_date"]
            if self.doc_metadata.get("doc_version"):
                base_meta["doc_version"] = self.doc_metadata["doc_version"]

        if chunk.source_type == "laws":
            payload = {
                "source_type": "laws",
                "article_num": chunk.article_number,
                "referenced_laws": chunk.referenced_laws or [],
                "content": chunk.content[:EMBEDDING_PAYLOAD_MAX_CHARS] if chunk.content else chunk.full_text[:EMBEDDING_PAYLOAD_MAX_CHARS],
                "type": "law_article",
            }
            # 注入层级信息和 YAML frontmatter 元数据
            if chunk.metadata:
                for key in ("title", "region", "category", "status", "effective_date",
                             "issuer", "doc_type", "doc_version", "h1", "h2", "bian", "zhang", "jie"):
                    if chunk.metadata.get(key):
                        payload[key] = chunk.metadata[key]
            # 文件级 doc_metadata 作为 fallback（覆盖优先级低于 chunk.metadata）
            payload.update(base_meta)
            return payload
        elif chunk.source_type == "standards":
            payload = {
                "title": chunk.chunk_id,
                "source_type": "standards",
                "category": chunk.category,
                "violation": chunk.violation,
                "legal_basis": chunk.legal_basis,
                "content": chunk.full_text[:EMBEDDING_PAYLOAD_MAX_CHARS],
            }
            if chunk.severity_levels:
                payload["severity_levels"] = [s.level for s in chunk.severity_levels]
                payload["penalty_summary"] = ";".join(s.penalty[:50] for s in chunk.severity_levels)
                payload["type"] = "discretion_standard"
            else:
                payload["applicable_conditions"] = chunk.applicable_conditions[:500]
                payload["list_type"] = chunk.list_type
                payload["violated_provision"] = chunk.violated_provision
                payload["penalty_basis"] = chunk.penalty_basis
                payload["type"] = "exemption_list"
            payload.update(base_meta)
            return payload
        elif chunk.source_type == "cases":
            payload = {
                "title": chunk.chunk_id,
                "source_type": "cases",
                "case_type": chunk.case_type,
                "violation_type": chunk.violation_type,
                "violation_category": chunk.violation_category,
                "legal_basis": chunk.legal_basis_list or [],
                "penalty_summary": chunk.penalty_summary,
                "discretion_level": chunk.discretion_level,
                "content": chunk.full_text[:EMBEDDING_PAYLOAD_MAX_CHARS],
                "type": "case_precedent",
            }
            payload.update(base_meta)
            return payload
        else:
            payload = {
                "title": chunk.chunk_id,
                "source_type": "other",
                "content": chunk.content[:EMBEDDING_PAYLOAD_MAX_CHARS] if chunk.content else chunk.full_text[:EMBEDDING_PAYLOAD_MAX_CHARS],
                "type": "md_document",
            }
            payload.update(base_meta)
            return payload

    async def _get_embedding(self, text: str) -> Optional[List[float]]:
        await self._ensure_config()
        for attempt in range(EMBEDDING_MAX_RETRIES):
            try:
                async with httpx.AsyncClient(timeout=30) as client:
                    res = await client.post(
                        f"{self.embed_api_base}/embeddings",
                        headers={
                            "Authorization": f"Bearer {self.embed_api_key}",
                            "Content-Type": "application/json",
                        },
                        json={"model": self.embed_model, "input": [text], "encoding_format": "float"},
                    )
                    if res.status_code == 200:
                        return res.json()["data"][0]["embedding"]
                    elif res.status_code == 429:
                        wait = EMBEDDING_RETRY_BACKOFF ** attempt
                        logger.warning(f"Embedding rate limited, retrying in {wait}s")
                        await asyncio.sleep(wait)
                    else:
                        logger.error(f"Embedding API error: {res.status_code} {res.text[:200]}")
                        return None
            except Exception as e:
                logger.error(f"Embedding API exception: {e}")
                if attempt < EMBEDDING_MAX_RETRIES - 1:
                    await asyncio.sleep(EMBEDDING_RETRY_BACKOFF ** attempt)
        return None

    async def _get_embeddings_batch(self, texts: List[str]) -> Optional[List[List[float]]]:
        await self._ensure_config()
        for attempt in range(EMBEDDING_MAX_RETRIES):
            try:
                async with httpx.AsyncClient(timeout=60) as client:
                    res = await client.post(
                        f"{self.embed_api_base}/embeddings",
                        headers={
                            "Authorization": f"Bearer {self.embed_api_key}",
                            "Content-Type": "application/json",
                        },
                        json={"model": self.embed_model, "input": texts, "encoding_format": "float"},
                    )
                    if res.status_code == 200:
                        data = res.json()["data"]
                        data.sort(key=lambda x: x["index"])
                        return [d["embedding"] for d in data]
                    elif res.status_code == 429:
                        wait = EMBEDDING_RETRY_BACKOFF ** attempt
                        logger.warning(f"Embedding batch rate limited, retrying in {wait}s")
                        await asyncio.sleep(wait)
                    else:
                        logger.error(f"Embedding batch API error: {res.status_code} {res.text[:200]}")
                        return None
            except Exception as e:
                logger.error(f"Embedding batch API exception: {e}")
                if attempt < EMBEDDING_MAX_RETRIES - 1:
                    await asyncio.sleep(EMBEDDING_RETRY_BACKOFF ** attempt)
        return None


class GraphitiChunkImporter:
    def __init__(self):
        self.graphiti_url = getattr(settings, 'GRAPHITI_BASE_URL', 'http://localhost:2082')

    async def import_chunk(self, chunk: DocumentChunk) -> dict:
        try:
            async with httpx.AsyncClient(timeout=GRAPHITI_TIMEOUT_SEC) as client:
                res = await client.post(
                    f"{self.graphiti_url}/api/knowledge-graph/add-knowledge-episode",
                    json={
                        "content": chunk.full_text[:4000],
                        "source": f"chunk_{chunk.chunk_id}",
                        "entity_types": self._get_entity_types(chunk.source_type),
                    },
                )
                if res.status_code in (200, 201):
                    data = res.json()
                    return {
                        "status": "done",
                        "entities_count": data.get("entities_count", 0),
                        "relations_count": data.get("relations_count", 0),
                    }
                else:
                    return {"status": "failed", "error": res.text[:300], "chunk_id": chunk.chunk_id}
        except Exception as e:
            return {"status": "failed", "error": str(e), "chunk_id": chunk.chunk_id}

    async def import_chunks_batch(self, chunks: List[DocumentChunk], concurrency: int = GRAPHITI_CONCURRENCY) -> dict:
        total = len(chunks)
        success_count = 0
        failed_count = 0
        total_entities = 0
        total_relations = 0
        errors = []

        sem = asyncio.Semaphore(concurrency)

        async def _process_one(chunk):
            nonlocal success_count, failed_count, total_entities, total_relations
            async with sem:
                result = await self.import_chunk(chunk)
                if result.get("status") == "done":
                    success_count += 1
                    total_entities += result.get("entities_count", 0)
                    total_relations += result.get("relations_count", 0)
                else:
                    failed_count += 1
                    if result.get("error"):
                        errors.append(f"{chunk.chunk_id}: {result['error'][:100]}")
                return result

        tasks = [_process_one(c) for c in chunks]
        await asyncio.gather(*tasks, return_exceptions=True)

        return {
            "status": "done" if failed_count == 0 else "partial",
            "total": total,
            "success_count": success_count,
            "failed_count": failed_count,
            "total_entities": total_entities,
            "total_relations": total_relations,
            "errors": errors[:10],
        }

    def _get_entity_types(self, source_type: str) -> list:
        if source_type == "laws":
            return ["Law", "Article", "Concept"]
        elif source_type == "standards":
            return ["Law", "Article", "Violation", "Penalty", "DiscretionStandard"]
        elif source_type == "cases":
            return ["Case", "Violation", "Penalty", "Law", "Organization"]
        return ["Law", "Article", "Concept", "Violation", "Penalty"]


class DocumentImporter:
    def __init__(self):
        self.chunk_engine = ChunkEngine()
        self.qdrant_importer = QdrantChunkImporter()
        self.graphiti_importer = GraphitiChunkImporter()

    async def import_document(
        self,
        doc_id: int,
        content: str,
        doc_name: str,
        category: str,
        options: ImportOptions,
        db=None,
    ) -> ImportResult:
        start_time = time.time()
        errors = []

        if options.chunk_strategy == "auto":
            chunk_content = options.csv_content if options.csv_content else content
            strategy = self.chunk_engine.auto_detect_strategy(chunk_content, category)
        else:
            strategy = options.chunk_strategy

        chunk_content = options.csv_content if options.csv_content else content
        chunk_result = self.chunk_engine.chunk(chunk_content, strategy, category, {
            "doc_id": doc_id,
            "doc_name": doc_name,
            "category": category,
        })

        self.qdrant_importer.doc_metadata = {}
        if options.region:
            self.qdrant_importer.doc_metadata["region"] = options.region
        if options.issuer:
            self.qdrant_importer.doc_metadata["issuer"] = options.issuer
        if options.effective_date:
            self.qdrant_importer.doc_metadata["effective_date"] = options.effective_date
        if options.doc_version:
            self.qdrant_importer.doc_metadata["doc_version"] = options.doc_version

        chunk_info = {
            "strategy": strategy,
            "total_chunks": chunk_result.total_chunks,
            "strategy_display_name": chunk_result.strategy_display_name,
        }

        wiki_result = None
        qdrant_result = None
        graphiti_result = None

        if options.import_qdrant and options.target_collection:
            qdrant_result = await self.qdrant_importer.import_chunks_batch(
                chunk_result.chunks, options.target_collection
            )
            if db and qdrant_result.get("point_ids"):
                conn, cur = db
                await cur.execute(
                    """UPDATE kb_md_documents SET qdrant_collection = %s,
                       qdrant_point_ids = %s WHERE id = %s""",
                    (options.target_collection, json.dumps(qdrant_result["point_ids"]), doc_id),
                )
            if qdrant_result.get("failed_count", 0) > 0:
                errors.extend(qdrant_result.get("errors", []))

        if options.import_graphiti:
            graphiti_result = await self.graphiti_importer.import_chunks_batch(
                chunk_result.chunks, options.concurrency
            )
            if db and graphiti_result.get("success_count", 0) > 0:
                conn, cur = db
                await cur.execute(
                    "UPDATE kb_md_documents SET graphiti_imported = 1 WHERE id = %s",
                    (doc_id,),
                )
            if graphiti_result.get("failed_count", 0) > 0:
                errors.extend(graphiti_result.get("errors", []))

        if db:
            conn, cur = db
            any_success = (
                (qdrant_result and qdrant_result.get("success_count", 0) > 0)
                or (graphiti_result and graphiti_result.get("success_count", 0) > 0)
            )
            if any_success:
                await cur.execute(
                    """UPDATE kb_md_documents SET import_status = 'imported',
                       chunk_strategy = %s, chunk_count = %s, updated_at = NOW()
                       WHERE id = %s""",
                    (strategy, chunk_result.total_chunks, doc_id),
                )

        elapsed = time.time() - start_time
        logger.info(f"Document import completed: doc_id={doc_id}, strategy={strategy}, "
                     f"chunks={chunk_result.total_chunks}, elapsed={elapsed:.1f}s")

        return ImportResult(
            success=any_success,
            chunk_info=chunk_info,
            qdrant_result=qdrant_result,
            graphiti_result=graphiti_result,
            errors=errors,
        )


document_importer = DocumentImporter()
