# owl_loader.py
from rdflib import Graph, Namespace, RDF, RDFS, OWL

# 온톨로지 파일 로드
g = Graph()
g.parse("investigation.ttl", format="turtle")

print(f"✅ 온톨로지 로드 완료 — 총 트리플(Triple) 수: {len(g)}")
print()

# ─── 예시 1: 모든 OWL 클래스 출력 ───
print("=== 정의된 클래스 목록 ===")
for cls in g.subjects(RDF.type, OWL.Class):
    local = str(cls).split("#")[-1]
    print(f"  클래스: {local}")

# ─── 예시 2: 상속 관계(subClassOf) 조회 ───
print("\n=== 상속(subClassOf) 관계 ===")
for child, _, parent in g.triples((None, RDFS.subClassOf, None)):
    child_name  = str(child).split("#")[-1]
    parent_name = str(parent).split("#")[-1]
    print(f"  {child_name}  →  subClassOf  →  {parent_name}")

# ─── 예시 3: disjoint 관계 확인 ───
print("\n=== Disjoint(상호 배타) 관계 ===")
for s, _, o in g.triples((None, OWL.disjointWith, None)):
    s_name = str(s).split("#")[-1]
    o_name = str(o).split("#")[-1]
    print(f"  {s_name}  ⊕  {o_name}  (동시에 속할 수 없음)")

# ─── 예시 4: ObjectProperty(관계) 목록 ───
print("\n=== ObjectProperty (관계 정의) ===")
for prop in g.subjects(RDF.type, OWL.ObjectProperty):
    prop_name = str(prop).split("#")[-1]
    domains = list(g.objects(prop, RDFS.domain))
    ranges  = list(g.objects(prop, RDFS.range))
    d = str(domains[0]).split("#")[-1] if domains else "?"
    r = str(ranges[0]).split("#")[-1] if ranges else "?"
    print(f"  :{prop_name}  ({d} → {r})")

# database.py
from neo4j import AsyncGraphDatabase, AsyncDriver

NEO4J_URI  = "bolt://localhost:7687"
NEO4J_USER = "neo4j"
NEO4J_PASS = "password1234!"

def create_driver() -> AsyncDriver:
    return AsyncGraphDatabase.driver(
        NEO4J_URI,
        auth=(NEO4J_USER, NEO4J_PASS),
        max_connection_pool_size=50,
    )

# models.py
from pydantic import BaseModel, Field, field_validator
from typing import Optional, List
from enum import Enum

class PersonRole(str, Enum):
    SUSPECT      = "Suspect"
    VICTIM       = "Victim"
    INVESTIGATOR = "Investigator"

class PersonCreate(BaseModel):
    name:         str             = Field(..., min_length=1, max_length=50)
    role:         PersonRole      = Field(...)
    age:          Optional[int]   = Field(None, ge=0, le=150)
    rank:         Optional[str]   = None
    risk_score:   Optional[float] = Field(None, ge=0.0, le=10.0)
    organization: Optional[str]   = None

class PersonResponse(BaseModel):
    name:       str
    role:       str
    age:        Optional[int]   = None
    rank:       Optional[str]   = None
    risk_score: Optional[float] = None

class OrganizationMember(BaseModel):
    name: str
    role: str
    age:  Optional[int] = None

class RelationCreate(BaseModel):
    from_name:         str            = Field(...)
    to_name:           str            = Field(...)
    relationship_type: str            = Field(...)
    metadata:          Optional[dict] = Field(default_factory=dict)

class PathResult(BaseModel):
    source:     str
    target:     str
    hops:       int
    path_nodes: List[str]

# main.py
from contextlib import asynccontextmanager
from typing import List
from fastapi import FastAPI, HTTPException, Depends, Request
from neo4j import AsyncSession

from database import create_driver
from models import (
    PersonCreate, PersonResponse,
    OrganizationMember, RelationCreate, PathResult
)

VALID_ROLES     = {"Suspect", "Victim", "Investigator"}
VALID_RELATIONS = {"KNOWS", "PARTY_TO", "INVESTIGATES"}

# ── 라이프스팬: 앱 시작/종료 시 드라이버 관리 ──
@asynccontextmanager
async def lifespan(app: FastAPI):
    driver = create_driver()
    app.state.driver = driver
    print("✅ Neo4j 드라이버 연결 완료")
    yield
    await driver.close()
    print("🔒 Neo4j 드라이버 종료 완료")

app = FastAPI(
    title="수사 인텔리전스 시스템",
    description="OWL 온톨로지 기반 수사 프로파일링 REST API",
    version="1.0.0",
    lifespan=lifespan,
)

# ── 세션 의존성 주입 ──
async def get_session(request: Request) -> AsyncSession:
    async with request.app.state.driver.session() as session:
        yield session

# ── /setup-sample-data ──
@app.post("/setup-sample-data", tags=["관리"])
async def setup_sample_data(session: AsyncSession = Depends(get_session)):
    await session.run("MATCH (n) DETACH DELETE n")
    data_query = """
    CREATE (o1:CriminalOrganization:Organization {name:'BlackShark', foundedYear:2010})
    CREATE (o2:Organization {name:'SeoulMetroPolice'})
    CREATE (p1:Person:Suspect {name:'강두식', age:45, riskScore:8.5})
    CREATE (p2:Person:Suspect {name:'이칠성', age:38, riskScore:7.2})
    CREATE (p3:Person:Victim  {name:'최민준', age:29})
    CREATE (p4:Person:Investigator {name:'박문수', age:32, rank:'경감'})
    CREATE (p5:Person:Investigator {name:'김정도', age:41, rank:'경정'})
    CREATE (c1:Crime:Event {id:'C-001', type:'마약 밀매',   status:'수사중'})
    CREATE (c2:DrugCrime:Crime:Event {id:'C-002', type:'필로폰 유통', status:'기소'})
    CREATE (p1)-[:MEMBER_OF {rank:'두목'}]->(o1)
    CREATE (p2)-[:MEMBER_OF {rank:'조직원'}]->(o1)
    CREATE (p4)-[:MEMBER_OF]->(o2)
    CREATE (p5)-[:MEMBER_OF]->(o2)
    CREATE (p1)-[:KNOWS {relationshipType:'동생'}]->(p2)
    CREATE (p2)-[:KNOWS]->(p3)
    CREATE (p1)-[:PARTY_TO {role:'주범', confidence:0.95}]->(c1)
    CREATE (p2)-[:PARTY_TO {role:'공범', confidence:0.88}]->(c1)
    CREATE (p3)-[:PARTY_TO {role:'피해자'}]->(c1)
    CREATE (p1)-[:PARTY_TO {role:'주범', confidence:0.99}]->(c2)
    CREATE (p4)-[:INVESTIGATES]->(c1)
    CREATE (p5)-[:INVESTIGATES]->(c2)
    CREATE (c1)-[:PART_OF]->(c2)
    """
    await session.run(data_query)
    return {"status": "success", "message": "샘플 데이터 생성 완료"}

# ── /persons/register ──
@app.post("/persons/register", tags=["인물"])
async def register_person(
    person: PersonCreate,
    session: AsyncSession = Depends(get_session),
):
    if person.role.value not in VALID_ROLES:
        raise HTTPException(status_code=400, detail=f"허용 역할: {VALID_ROLES}")
    role_label = person.role.value
    query = f"""
    MERGE (p:Person:{role_label} {{name: $name}})
    SET p.age       = $age,
        p.riskScore = $risk_score,
        p.rank      = $rank
    WITH p
    FOREACH (_ IN CASE WHEN $org IS NOT NULL THEN [1] ELSE [] END |
        MERGE (o:Organization {{name: $org}})
        MERGE (p)-[:MEMBER_OF]->(o)
    )
    RETURN p.name AS name
    """
    try:
        result = await session.run(
            query,
            name=person.name, age=person.age,
            risk_score=person.risk_score, rank=person.rank,
            org=person.organization,
        )
        record = await result.single()
        return {"status": "success", "message": f"'{person.name}' 등록 완료"}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

# ── /persons/{name} ──
@app.get("/persons/{name}", response_model=PersonResponse, tags=["인물"])
async def get_person(name: str, session: AsyncSession = Depends(get_session)):
    query = """
    MATCH (p:Person {name: $name})
    RETURN p.name AS name, p.age AS age, p.rank AS rank,
           p.riskScore AS risk_score,
           [lbl IN labels(p) WHERE lbl <> 'Person'][0] AS role
    """
    result = await session.run(query, name=name)
    record = await result.single()
    if not record:
        raise HTTPException(status_code=404, detail=f"'{name}' 인물 없음")
    return PersonResponse(
        name=record["name"], role=record["role"] or "Unknown",
        age=record["age"], rank=record["rank"], risk_score=record["risk_score"],
    )

# ── /suspects ──
@app.get("/suspects", response_model=List[PersonResponse], tags=["인물"])
async def list_suspects(session: AsyncSession = Depends(get_session)):
    query = """
    MATCH (s:Suspect)
    RETURN s.name AS name, s.age AS age, s.riskScore AS risk_score, 'Suspect' AS role
    ORDER BY s.riskScore DESC
    """
    result = await session.run(query)
    return [
        PersonResponse(name=r["name"], role=r["role"], age=r["age"], risk_score=r["risk_score"])
        async for r in result
    ]

# ── /organizations/{org_name}/members ──
@app.get("/organizations/{org_name}/members",
         response_model=List[OrganizationMember], tags=["조직"])
async def get_organization_chart(org_name: str, session: AsyncSession = Depends(get_session)):
    query = """
    MATCH (p:Person)-[:MEMBER_OF]->(o:Organization {name: $org_name})
    RETURN p.name AS name, p.age AS age,
           [lbl IN labels(p) WHERE lbl <> 'Person'][0] AS role
    ORDER BY p.name
    """
    result = await session.run(query, org_name=org_name)
    members = [
        OrganizationMember(name=r["name"], role=r["role"] or "Unknown", age=r["age"])
        async for r in result
    ]
    if not members:
        raise HTTPException(status_code=404, detail=f"'{org_name}' 조직 없음")
    return members

# ── /analysis/path ──
@app.get("/analysis/path", tags=["분석"])
async def find_path(
    from_name: str, to_name: str,
    max_hops: int = 6,
    session: AsyncSession = Depends(get_session),
):
    if not (1 <= max_hops <= 10):
        raise HTTPException(status_code=400, detail="max_hops는 1~10 사이")
    query = f"""
    MATCH (a:Person {{name: $from_name}}), (b:Person {{name: $to_name}})
    MATCH path = shortestPath((a)-[*..{max_hops}]-(b))
    RETURN [n IN nodes(path) | coalesce(n.name, n.id, 'unknown')] AS path_nodes,
           length(path) AS hops
    LIMIT 1
    """
    result = await session.run(query, from_name=from_name, to_name=to_name)
    record = await result.single()
    if not record:
        return {"status": "not_found", "message": f"{max_hops}단계 이내 경로 없음"}
    return {"status": "found", "from": from_name, "to": to_name,
            "hops": record["hops"], "path_nodes": record["path_nodes"]}

# ── /analysis/integrity ──
@app.get("/analysis/integrity", tags=["분석"])
async def check_integrity(session: AsyncSession = Depends(get_session)):
    query = """
    MATCH (n:Investigator) WHERE n:Suspect
    RETURN n.name AS name, labels(n) AS labels, '논리 위반 (disjointWith)' AS violation
    """
    result = await session.run(query)
    violations = [
        {"name": r["name"], "labels": r["labels"], "violation": r["violation"]}
        async for r in result
    ]
    return {"status": "ok" if not violations else "violations_found",
            "violations": violations, "count": len(violations)}