from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Literal, Optional
from uuid import uuid4
import ast

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field
import uvicorn
import traceback
from loguru import logger


from utils.common_util import extract_url_file, format_now
from utils.http_util import download_file
from core.cache import get_cached_doc_tool, get_cached_memory
from core.config import (
    doc_support_formats,
    pdf_support_formats,
    MERGE_RULE_PROMPT,
    use_lufa,
)
from core.tools.segment_summary import SegmentSummaryTool
from core.tools.segment_review import SegmentReviewTool
from core.tools.segment_rule_router import SegmentRuleRouterTool
from core.tools.rule_filter import LufaPartyRuleFilterTool
from core.tools.retrieve_reference import RetrieveReferenceTool
from core.tools.reflect_retry import ReflectRetryTool
from core.tools.segment_merger import SegmentMergerTool
from core.tools.fact_merger import FactMergerTool
from core.memory import Finding
from core.memory import FINDING_KEY_MERGE, FINDING_KEY_REFLECT, FINDING_KEY_REVIEW

app = FastAPI(title="合同审查智能体", version="0.1.0")
TMP_DIR = Path(__file__).resolve().parent / "tmp"
TMP_DIR.mkdir(parents=True, exist_ok=True)
summary_tool = SegmentSummaryTool()
review_tool = SegmentReviewTool()
rule_router_tool = SegmentRuleRouterTool()
lufa_party_rule_filter_tool = LufaPartyRuleFilterTool()
reference_tool = RetrieveReferenceTool()
reflect_tool = ReflectRetryTool()
merger_tool = SegmentMergerTool()
fact_merger_tool = FactMergerTool()


@app.post("/sleep")
def sleep(t: int):
    import time

    time.sleep(t)
    return {"res": f"sleep over for {t} seconds."}


########################################################################################################################


class DocumentParseRequest(BaseModel):
    conversation_id: str
    urls: List[str] = Field(..., description="File download url")
    file_ext: Optional[str] = None
    ruleset_id: Optional[str] = "通用"


class DocumentParseResponse(BaseModel):
    conversation_id: str
    segment_ids: List[int]
    ruleset_items: List[str]
    summary_names: List[str]
    text: Optional[str] = None
    file_ext: Optional[str] = None
    file_name: Optional[str] = None


@app.post("/documents/parse", response_model=DocumentParseResponse)
async def parse_document(payload: DocumentParseRequest) -> DocumentParseResponse:
    if not payload.urls:
        raise HTTPException(status_code=400, detail="No URLs provided")
    try:
        file_path = download_file(payload.urls[0], TMP_DIR)
        if not file_path:
            raise RuntimeError("download returned empty path")
    except Exception as exc:
        raise HTTPException(status_code=500, detail=f"Download failed: {exc}")
    # get doc tool
    file_ext = payload.file_ext or Path(file_path).suffix
    try:
        doc_obj, _ = get_cached_doc_tool(payload.conversation_id, file_ext)
    except Exception as exc:
        raise HTTPException(
            status_code=400, detail=f"Document tool not available: {exc}"
        )
    doc_obj.load(file_path)
    # ocr
    await doc_obj.get_from_ocr()
    text = doc_obj.get_all_text()
    segment_ids = doc_obj.get_chunk_id_list()
    # TODO: FastGPT BUG segment_ids必须从1开始，0开始会缺少第一段文本，后续需要修复
    segment_ids = [idx + 1 for idx in segment_ids]
    # get ruleset items
    ruleset_id = payload.ruleset_id or reference_tool.default_ruleset_id
    ruleset_items = reference_tool.run(ruleset_id=ruleset_id).get("rules", [])
    ruleset_review_items = [
        t
        for t in (r.get("title") for r in ruleset_items)
        if isinstance(t, str) and t.strip()
    ]
    summary_names = list(
        dict.fromkeys(
            s.strip()
            for s in (r.get("summary") for r in ruleset_items)
            if isinstance(s, str) and s.strip()
        )
    )
    return DocumentParseResponse(
        conversation_id=payload.conversation_id,
        text=text,
        segment_ids=segment_ids,
        ruleset_items=ruleset_review_items,
        summary_names=summary_names,
        file_ext=file_ext,
        file_name=Path(file_path).name,
    )


########################################################################################################################


class SegmentSummaryRequest(BaseModel):
    conversation_id: str
    segment_id: int
    party_role: Optional[str] = ""
    ruleset_id: Optional[str] = "通用"
    routed_rule_titles: Optional[List[str]] = None
    routed_summary_names: Optional[List[str]] = None
    file_ext: str
    context_facts: Optional[Dict] = None


class SegmentSummaryResponse(BaseModel):
    conversation_id: str
    segment_id: int
    facts: Dict


@app.post("/segments/summary/facts", response_model=SegmentSummaryResponse)
def summarize_facts(payload: SegmentSummaryRequest) -> SegmentSummaryResponse:
    store = get_cached_memory(payload.conversation_id)
    try:
        doc_obj, _ = get_cached_doc_tool(payload.conversation_id, payload.file_ext)
    except Exception as exc:
        raise HTTPException(
            status_code=400, detail=f"Document tool not available: {exc}"
        )

    segment_idx = payload.segment_id - 1
    try:
        segment_text = doc_obj.get_chunk_item(segment_idx)
    except Exception as exc:
        raise HTTPException(
            status_code=404,
            detail=f"Segment text not found for id {payload.segment_id}: {exc}. Please parse document first.",
        )

    ruleset_id = payload.ruleset_id or reference_tool.default_ruleset_id
    if payload.routed_summary_names is not None:
        summary_names = {
            name.strip()
            for name in payload.routed_summary_names
            if isinstance(name, str) and name.strip()
        }
        all_rules = reference_tool.run(ruleset_id=ruleset_id).get("rules", [])
        rules = [
            rule
            for rule in all_rules
            if str(rule.get("summary", "")).strip() in summary_names
        ]
    else:
        rules = reference_tool.run(
            ruleset_id=ruleset_id,
            routed_rule_titles=payload.routed_rule_titles,
        ).get("rules", [])
    result = summary_tool.run(
        segment_id=segment_idx,
        segment_text=segment_text,
        rules=rules,
        party_role=payload.party_role or "",
        context_facts=payload.context_facts,
    )
    store.add_facts(result)

    return SegmentSummaryResponse(
        conversation_id=payload.conversation_id,
        segment_id=payload.segment_id,
        facts=result,
    )


########################################################################################################################


class SegmentReviewRequest(BaseModel):
    conversation_id: str
    segment_id: int
    party_role: Optional[str] = ""
    ruleset_id: Optional[str] = "通用"
    routed_rule_titles: Optional[List[str]] = None
    file_ext: str
    context_memories: Optional[List[Dict]] = None
    route_by: Literal["rule", "summary"] = Field(
        default="rule", description="路由依据：rule=审查规则项，summary=摘要项"
    )


class SegmentReviewResponse(BaseModel):
    conversation_id: str
    segment_id: int
    overall_conclusion: str
    findings: List[Dict]


class SegmentRuleRouterResponse(BaseModel):
    conversation_id: str
    segment_id: int
    ruleset_id: str
    route_by: Literal["rule", "summary"] = "rule"
    routed_rule_titles: List[str]
    routed_summary_names: List[str] = Field(default_factory=list)
    routed_rules: List[Dict]


@app.post("/segments/review/findings", response_model=SegmentReviewResponse)
def review_segment(payload: SegmentReviewRequest) -> SegmentReviewResponse:
    store = get_cached_memory(payload.conversation_id)
    try:
        doc_obj, _ = get_cached_doc_tool(payload.conversation_id, payload.file_ext)
    except Exception as exc:
        raise HTTPException(
            status_code=400, detail=f"Document tool not available: {exc}"
        )

    segment_idx = payload.segment_id - 1
    try:
        segment_text = doc_obj.get_chunk_item(segment_idx)
    except Exception as exc:
        raise HTTPException(
            status_code=404,
            detail=f"Segment text not found for id {payload.segment_id}: {exc}. Please parse document first.",
        )

    ruleset_id = payload.ruleset_id or reference_tool.default_ruleset_id
    rules = reference_tool.run(
        ruleset_id=ruleset_id,
        routed_rule_titles=payload.routed_rule_titles,
    ).get("rules", [])
    # 暂时不添加摘要看下结果
    # summary_keywords = reference_tool.summary_keywords(rules)
    # context_summaries = store.search_facts(summary_keywords)

    result = review_tool.run(
        segment_id=segment_idx,
        segment_text=segment_text,
        rules=rules,
        party_role=payload.party_role or "",
        # 暂时不添加摘要看下结果
        # context_summaries=context_summaries,
        # TODO 获取与当前审查相关的上下文记忆（如之前的审查结果、总结事实等），而非全部记忆
        context_memories=payload.context_memories,
        merge_rules_prompt=MERGE_RULE_PROMPT,  # TODO 是否合并规则到同一 prompt 进行审查，当前默认合并，后续可调整为不合并以提升审查的针对性
    )

    # Persist findings to memory store
    for f in result.get("findings", []) or []:
        try:
            store.add_finding(
                FINDING_KEY_REVIEW,
                Finding.from_dict(
                    {
                        "rule_title": f.get("rule_title", ""),
                        "segment_id": segment_idx,
                        "original_text": f.get("original_text", ""),
                        "issue": f.get("issue", ""),
                        "risk_level": (
                            f.get("risk_level") or f.get("level") or ""
                        ).upper(),
                        "result": f.get("result", ""),
                        "suggestion": f.get("suggestion", ""),
                    }
                ),
            )
        except Exception as e:
            logger.error(e)
            continue
    return SegmentReviewResponse(
        conversation_id=payload.conversation_id,
        segment_id=payload.segment_id,
        overall_conclusion=result.get("overall_conclusion", ""),
        findings=result.get("findings", []),
    )


@app.post("/segments/review/rule-router", response_model=SegmentRuleRouterResponse)
def route_segment_rules(payload: SegmentReviewRequest) -> SegmentRuleRouterResponse:
    try:
        doc_obj, _ = get_cached_doc_tool(payload.conversation_id, payload.file_ext)
    except Exception as exc:
        raise HTTPException(
            status_code=400, detail=f"Document tool not available: {exc}"
        )

    segment_idx = payload.segment_id - 1
    try:
        segment_text = doc_obj.get_chunk_item(segment_idx)
    except Exception as exc:
        raise HTTPException(
            status_code=404,
            detail=f"Segment text not found for id {payload.segment_id}: {exc}. Please parse document first.",
        )

    ruleset_id = payload.ruleset_id or reference_tool.default_ruleset_id
    rules = reference_tool.run(ruleset_id=ruleset_id).get("rules", [])

    if use_lufa and rules:
        try:
            total_segments = len(doc_obj.get_chunk_id_list() or [])
        except Exception:
            total_segments = 0

        filtered_payload = lufa_party_rule_filter_tool.run(
            {
                "rules": rules,
                "segment_idx": segment_idx,
                "total_segments": total_segments,
            }
        )
        rules = filtered_payload.get("rules", rules)

    result = rule_router_tool.run(
        segment_id=segment_idx,
        segment_text=segment_text,
        rules=rules,
        party_role=payload.party_role or "",
        context_memories=payload.context_memories,
        route_by=payload.route_by,
    )

    return SegmentRuleRouterResponse(
        conversation_id=payload.conversation_id,
        segment_id=payload.segment_id,
        ruleset_id=ruleset_id,
        route_by=result.get("route_by", payload.route_by),
        routed_rule_titles=result.get("routed_rule_titles", []),
        routed_summary_names=result.get("routed_summary_names", []),
        routed_rules=result.get("routed_rules", []),
    )


########################################################################################################################
class ReflectReviewRequest(BaseModel):
    conversation_id: str
    party_role: str
    ruleset_id: Optional[str] = "通用"
    rule_title: str


class ReflectReviewResponse(BaseModel):
    conversation_id: str
    rule_title: str
    findings: List[Dict]


class MergerRequest(BaseModel):
    conversation_id: str
    segment_id: int
    finding_key: Optional[str] = FINDING_KEY_REVIEW


class MergerResponse(BaseModel):
    conversation_id: str
    segment_id: int
    finding_key: str
    merged_findings: List[Dict]


class FactsMergerRequest(BaseModel):
    conversation_id: str
    summary_names: List[str]
    merge_mode: Literal["llm", "rule"] = "llm"


class FactsMergerResponse(BaseModel):
    conversation_id: str
    summary_names: List[str]
    merge_facts: Dict


@app.post("/segments/review/reflect", response_model=ReflectReviewResponse)
def reflect_review(payload: ReflectReviewRequest) -> ReflectReviewResponse:
    store = get_cached_memory(payload.conversation_id)
    ruleset_id = payload.ruleset_id or reference_tool.default_ruleset_id
    ruleset_items = reference_tool.run(ruleset_id=ruleset_id).get("rules", [])
    rule = next(
        (r for r in ruleset_items if r.get("title") == payload.rule_title), None
    )
    if not rule:
        raise HTTPException(
            status_code=404, detail=f"Rule not found: {payload.rule_title}"
        )
    summary_keywords = reference_tool.summary_keywords([rule])
    context_summaries_facts = store.search_facts(summary_keywords)
    # 查找审查规则对应的 findings
    findings = [
        f.__dict__
        for f in store.search_findings(
            FINDING_KEY_REVIEW, "", rule_title=payload.rule_title
        )
    ]
    final_findings = reflect_tool.run(
        party_role=payload.party_role,
        rule=rule,
        facts=context_summaries_facts,
        findings=findings,
    )

    for f in final_findings or []:
        try:
            store.add_finding(
                FINDING_KEY_REFLECT,
                Finding.from_dict(
                    {
                        "rule_title": f.get("rule_title", ""),
                        "segment_id": f.get("segment_id", 0),
                        "original_text": f.get("original_text", ""),
                        "issue": f.get("issue", ""),
                        "risk_level": (
                            f.get("risk_level") or f.get("level") or "H"
                        ).upper(),
                        "suggestion": f.get("suggestion", ""),
                        "result": f.get("result", ""),
                    }
                ),
            )
        except Exception:
            continue
    return ReflectReviewResponse(
        conversation_id=payload.conversation_id,
        rule_title=payload.rule_title,
        findings=final_findings or [],
    )


@app.post("/segments/review/merger", response_model=MergerResponse)
def merge_segment_findings(payload: MergerRequest) -> MergerResponse:
    store = get_cached_memory(payload.conversation_id)
    source_key = payload.finding_key or FINDING_KEY_REVIEW
    target_segment_id = payload.segment_id - 1

    segment_findings = store.get_findings_by_segment(source_key, target_segment_id)
    unqualified_findings = [
        f for f in segment_findings if (f.result or "").strip() == "不合格"
    ]
    merged_result = merger_tool.run([f.__dict__ for f in unqualified_findings])
    merged_findings = merged_result.get("findings", []) or []

    for f in merged_findings:
        try:
            store.add_finding(
                FINDING_KEY_MERGE,
                Finding.from_dict(
                    {
                        "id": f.get("id", ""),
                        "rule_title": f.get("rule_title", ""),
                        "segment_id": target_segment_id,
                        "original_text": f.get("original_text", ""),
                        "issue": f.get("issue", ""),
                        "risk_level": (
                            f.get("risk_level") or f.get("level") or ""
                        ).upper(),
                        "suggestion": f.get("suggestion", ""),
                        "result": f.get("result", ""),
                    }
                ),
            )
        except Exception as e:
            logger.error(e)
            continue

    return MergerResponse(
        conversation_id=payload.conversation_id,
        segment_id=target_segment_id,
        finding_key=source_key,
        merged_findings=merged_findings,
    )


@app.post("/segments/summary/facts/merger", response_model=FactsMergerResponse)
def merge_summary_facts(payload: FactsMergerRequest) -> FactsMergerResponse:
    summary_names = [
        name.strip()
        for name in (payload.summary_names or [])
        if isinstance(name, str) and name.strip()
    ]
    if not summary_names:
        raise HTTPException(status_code=400, detail="summary_names cannot be empty")

    store = get_cached_memory(payload.conversation_id)
    result = fact_merger_tool.run(
        facts=store.get_facts(),
        summary_names=summary_names,
        merge_mode=payload.merge_mode,
    )
    merge_facts = result.get("merge_facts", {}) or {}
    if merge_facts:
        store.add_merge_facts(merge_facts)

    return FactsMergerResponse(
        conversation_id=payload.conversation_id,
        summary_names=summary_names,
        merge_facts=merge_facts,
    )


########################################################################################################################
class ConversationResponse(BaseModel):
    conversation_id: str
    created_at: str


@app.post("/conversations/new", response_model=ConversationResponse)
def new_conversation() -> ConversationResponse:
    conversation_id = uuid4().hex
    created_at = format_now()
    return ConversationResponse(conversation_id=conversation_id, created_at=created_at)


########################################################################################################################


class FactsRetrieveRequest(BaseModel):
    conversation_id: str
    keywords: List[str] = Field(..., description="facts 检索关键字列表")


class FactsRetrieveResponse(BaseModel):
    conversation_id: str
    keywords: List[str]
    facts: List[Any]
    total: int


@app.post("/memory/facts/retrieve", response_model=FactsRetrieveResponse)
def retrieve_facts(payload: FactsRetrieveRequest) -> FactsRetrieveResponse:
    keywords = [
        k.strip() for k in (payload.keywords or []) if isinstance(k, str) and k.strip()
    ]
    if not keywords:
        raise HTTPException(status_code=400, detail="keywords cannot be empty")

    store = get_cached_memory(payload.conversation_id)
    matched_facts = store.search_facts(keywords)
    return FactsRetrieveResponse(
        conversation_id=payload.conversation_id,
        keywords=keywords,
        facts=matched_facts,
        total=len(matched_facts),
    )


########################################################################################################################


class MemoryExportRequest(BaseModel):
    conversation_id: str
    file_ext: str
    file_name: Optional[str] = None
    finding_key: Optional[str] = FINDING_KEY_REVIEW


class MemoryExportResponse(BaseModel):
    conversation_id: str
    excel_url: str
    doc_url: str


class FactsMemoryExportRequest(BaseModel):
    conversation_id: str
    file_name: Optional[str] = None


class FactsMemoryExportResponse(BaseModel):
    conversation_id: str
    excel_url: Any


@app.post("/memory/export", response_model=MemoryExportResponse)
def export_memory(payload: MemoryExportRequest) -> MemoryExportResponse:
    store = get_cached_memory(payload.conversation_id)
    try:
        doc_obj, _ = get_cached_doc_tool(payload.conversation_id, payload.file_ext)
    except Exception as exc:
        raise HTTPException(
            status_code=400, detail=f"Document tool not available: {exc}"
        )
    try:
        excel_res = store.export_to_excel(file_name=payload.file_name)
    except ImportError as exc:
        raise HTTPException(status_code=500, detail=str(exc))
    except Exception as exc:
        raise HTTPException(status_code=500, detail=f"Export failed: {exc}")

    try:
        doc_res = store.export_findings_to_doc_comments(
            doc_obj,
            file_name=payload.file_name,
            finding_key=payload.finding_key or FINDING_KEY_REVIEW,
        )
    except Exception as exc:
        traceback.print_exc()
        raise HTTPException(
            status_code=500, detail=f"Export doc comments failed: {exc}"
        )

    return MemoryExportResponse(
        conversation_id=payload.conversation_id,
        excel_url=excel_res,
        doc_url=doc_res,
    )


@app.post("/memory/facts/export", response_model=FactsMemoryExportResponse)
def export_facts_memory(payload: FactsMemoryExportRequest) -> FactsMemoryExportResponse:
    store = get_cached_memory(payload.conversation_id)
    try:
        excel_res = store.export_facts_to_excel(file_name=payload.file_name)
    except ImportError as exc:
        raise HTTPException(status_code=500, detail=str(exc))
    except Exception as exc:
        raise HTTPException(status_code=500, detail=f"Export facts failed: {exc}")

    return FactsMemoryExportResponse(
        conversation_id=payload.conversation_id,
        excel_url=excel_res,
    )


if __name__ == "__main__":
    from core.config import use_lufa

    if use_lufa:
        port = 18168
    else:
        port = 18169
    uvicorn.run("main:app", host="0.0.0.0", port=port, log_level="info", reload=False)
