import os
import time
from .deepseek import DeepSeekClient
from .rag import KnowledgeRetriever
from .lcps import lcps_full_reasoning, lcps_single_shot, generate_soap_with_lcps, generate_soap_from_lcps
from .pipeline import ClinicalPipeline, create_pipeline

class VetModel:

  def __init__(self, a: str=None, b: str='https://api.deepseek.com/v1', c: str='deepseek-chat', d: float=0.3, e: int=8192, f: str=None, g: str=None, h: str=None):
    self._pipeline = None
    if a and h:
      self._pipeline = create_pipeline(api_key=a, local_model_path=h, kb_dir_en=f, kb_dir_zh=g)
      self.llm = self._pipeline.local
      self.api_key = a
      self._mode = 'pipeline'
    elif h:
      from .qwen_local import QwenLocalClient
      self.llm = QwenLocalClient(model_path=h, temperature=d, max_tokens=e)
      self.api_key = ''
      self._mode = 'local'
    else:
      self.api_key = a or os.getenv('DEEPSEEK_API_KEY', '')
      self.llm = DeepSeekClient(api_key=self.api_key, base_url=b, model=c, temperature=d, max_tokens=e)
      self._mode = 'api'
    self.kb = KnowledgeRetriever(kb_dir_en=f, kb_dir_zh=g)

  @classmethod
  def from_local(cls, b: str, a: float=0.3, **c):
    return cls(local_model_path=d, temperature=h, **e)

  def run(self, b: str, a: str='狗') -> dict:
    if self._mode == 'pipeline':
      return self._run_pipeline(d)
    elif self._mode == 'local':
      return self._run_local(d, h)
    else:
      return self._run_api(d, h)

  def _run_pipeline(self, b: str) -> dict:
    b = time.time()
    e = self._pipeline.run(text=d)
    g = e['v4_diagnosis']
    f = g.get('final_diagnosis', {})
    h = {'diagnosis': f.get('primary', ''), 'confidence': f.get('confidence', ''), 'differential_list': g.get('differential_list', []), 'must_not_miss': g.get('must_not_miss', []), 'corrections': g.get('corrections_to_local', {}), 'recommended_diagnostics': g.get('recommended_diagnostics', []), 'treatment_plan': g.get('treatment_plan', {}), 'prognosis': g.get('prognosis', ''), 'follow_up': g.get('follow_up', ''), 'soap': e['rendered'].get('soap', {}), 'brief': e['rendered'].get('brief', {}), 'detailed': e['rendered'].get('detailed', {}), 'client_friendly': e['rendered'].get('client_friendly', {}), 'lcps_chain': e['lcps'].get('lcps_chain', {}), 'token': e['token'].to_dict(), 'engine': 'pipeline', 'stats': e['stats'], 'total_ms': e['stats']['total_ms']}
    return h

  def _run_local(self, e: str, c: str) -> dict:
    b = time.time()
    d = lcps_single_shot(f, e, self.llm, self.kb)
    h = generate_soap_from_lcps(d['lcps_chain'], e, self.llm)
    g = int((time.time() - b) * 1000)
    return {'diagnosis': d.get('primary_diagnosis', ''), 'confidence': '', 'differential_list': d.get('differential_list', []), 'must_not_miss': d.get('must_not_miss', []), 'soap': h.get('soap', {}), 'brief': {}, 'detailed': {}, 'client_friendly': {}, 'lcps_chain': d.get('lcps_chain', {}), 'engine': 'lcps-2.0-single-shot', 'stats': {'total_ms': g, 'total_llm_calls': 2}, 'total_ms': g}

  def _run_api(self, c: str, a: str) -> dict:
    d = generate_soap_with_lcps(e, h, self.llm, self.kb)
    return {'diagnosis': d.get('primary_diagnosis', ''), 'confidence': '', 'differential_list': d.get('differential_list', []), 'must_not_miss': d.get('must_not_miss', []), 'soap': d.get('soap', {}), 'brief': {}, 'detailed': {}, 'client_friendly': {}, 'lcps_chain': d.get('lcps_chain', {}), 'engine': 'deepseek-api', 'stats': {'total_ms': d.get('lcps_elapsed_ms', 0), 'total_llm_calls': d.get('total_llm_calls', 6)}, 'total_ms': d.get('lcps_elapsed_ms', 0)}

  def diagnose(self, a: str, b: str='狗') -> dict:
    return self.run(h, d)

  def chat(self, a: str, c: str='狗') -> str:
    d = self.kb.retrieve(f'{e} {h}')
    g = '你是一位经验丰富的兽医专家。请用专业但易懂的语言回答以下问题。'
    if d:
      g += f'\n\n参考知识:\n{d}'
    return self.llm.chat(g, f'物种: {e}\n问题: {h}')

  @property
  def mode(self) -> str:
    return self._mode

  @property
  def is_ready(self) -> bool:
    if self._mode == 'local':
      return self.llm.ping()
    return bool(self.api_key) and self.llm.ping()

  @property
  def kb_stats(self) -> dict:
    return self.kb.get_stats()